Merge pull request #579 from avik-pal/patch-1

Update the activation inplace for cudnn
FluxML · May 1, 2024 · d0da256 · d0da256
2 parents 6d168b4 + 01e127e
commit d0da256
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ext/NNlibCUDACUDNNExt/conv.jl b/ext/NNlibCUDACUDNNExt/conv.jl
@@ -129,7 +129,7 @@ function conv_bias_act!(y::DenseCuArray{T}, x::DenseCuArray{T}, w::DenseCuArray{
     activation = (σ == NNlib.relu ? CUDNN_ACTIVATION_RELU : CUDNN_ACTIVATION_IDENTITY)
     cudnnConvolutionForward!(y, w, x, d; z, bias, activation, alpha, beta)
     if activation === CUDNN_ACTIVATION_IDENTITY && σ ∉ (nothing, identity)
-        y = σ.(y)
+        @. y = σ(y)
     end
     return y
 end