correct splatting of tail add first test for multiple dimensions

Co-Authored-By: Dhairya Gandhi <[email protected]>
FluxML · Jul 1, 2020 · d492c0f · d492c0f
1 parent 39b14fb
commit d492c0f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -154,7 +154,7 @@ outdims(m, (10, 2)) == (5, 2)
 """
 function outdims(l::Dense, isize)
     first(isize) == size(l.W, 2) || throw(DimensionMismatch("input size should equal to ($(size(l.W, 2)), ...), got $isize"))
-    return (size(l.W, 1), Base.tail(isize))
+    return (size(l.W, 1), Base.tail(isize)...)
 end
 
 """

diff --git a/test/layers/basic.jl b/test/layers/basic.jl
@@ -107,6 +107,7 @@ import Flux: activations
     m = Dense(10, 5)
     @test_throws DimensionMismatch Flux.outdims(m, (5, 2)) == (5,)
     @test Flux.outdims(m, (10,)) == (5,)
+    @test Flux.outdims(m, (10, 2)) == (5, 2)
 
     m = Chain(Dense(10, 8, σ), Dense(8, 5), Dense(5, 2))
     @test Flux.outdims(m, (10,)) == (2,)