Apply suggestion to check for d > 0 as well

Co-authored-by: Anton Smirnov <[email protected]>
FluxML · Jul 31, 2024 · 920f362 · 920f362
1 parent f933bff
commit 920f362
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -191,7 +191,7 @@ Dense(W::LinearAlgebra.Diagonal, bias = true, σ = identity) =
   Scale(W.diag, bias, σ)
 
 function _size_check(layer, x::AbstractArray, (d, n)::Pair)
-  d <= ndims(x) || throw(DimensionMismatch(string("layer ", layer,
+  0 < d <= ndims(x) || throw(DimensionMismatch(string("layer ", layer,
     " expects ndims(input) >= ", d, ", but got ", summary(x))))
   size(x, d) == n || throw(DimensionMismatch(string("layer ", layer,
     lazy" expects size(input, $d) == $n, but got ", summary(x))))