diff --git a/src/cuda/cudnn.jl b/src/cuda/cudnn.jl index 6d15fa61..81c2bcb4 100644 --- a/src/cuda/cudnn.jl +++ b/src/cuda/cudnn.jl @@ -179,8 +179,8 @@ end import ..Flux: Flux import ..Tracker: track, back, @back, istracked, TrackedArray -(BN::Flux.BatchNorm)(x::Union{CuParam{T,4},CuParam{T,5}}) where T<:Union{Float32, Float64} = - batchnorm(BN.γ, BN.β, x, BN.μ, BN.σ, BN.momentum; cache = nothing, alpha = 1, beta = 0, eps = BN.ϵ, training = BN.active) +(BN::Flux.BatchNorm)(x::Union{CuParam{T,4},CuParam{T,5}}, cache = nothing) where T<:Union{Float32, Float64} = + batchnorm(BN.γ, BN.β, x, BN.μ, BN.σ, BN.momentum; cache = cache, alpha = 1, beta = 0, eps = BN.ϵ, training = BN.active) _batchnorm(g, b, x, running_mean, running_var, momentum, cache, alpha, beta, eps, training) =