Skip conv_bias_act with Flux.Zeros

FluxML · Feb 7, 2022 · 8b5b26a · 8b5b26a
1 parent 6573f65
commit 8b5b26a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/layers/conv.jl b/src/layers/conv.jl
@@ -166,6 +166,7 @@ function (c::Conv)(x::AbstractArray)
 end
 
 _conv_bias_act(x, w, cdims, b, σ) = NNlib.conv_bias_act(x, w, cdims, b, σ)
+_conv_bias_act(x, w, cdims, ::Zeros, σ) = σ.(conv(x, w, cdims))
 _conv_bias_act(x::CuArray, w::CuArray, cdims, b::Zeros, σ) =
   _conv_bias_act(x, w, cdims, CUDA.zeros(size(b)...), σ)