Update src/layers/basic.jl

Co-authored-by: Brian Chen <[email protected]>
FluxML · Apr 2, 2022 · 7465575 · 7465575
1 parent 1f68819
commit 7465575
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -200,7 +200,7 @@ Diagonal(sz::Integer...; σ = identity, bias = true, init = ones32) = Diagonal(i
 
 function (a::Diagonal)(x::AbstractArray)
   σ = NNlib.fast_act(a.σ, x)  # replaces tanh => tanh_fast, etc
-  return σ == typeof(identity) ? a.scale .* x .+ a.bias : σ.(a.scale .* x .+ a.bias)
+  return σ === typeof(identity) ? a.scale .* x .+ a.bias : σ.(a.scale .* x .+ a.bias)
 end
 
 function Base.show(io::IO, l::Diagonal)