Work around GPU launch bug for BatchNorm
This commit is contained in:
parent
f0d5624ed2
commit
a740dadf6a
@ -138,7 +138,10 @@ function (BN::BatchNorm)(x)
|
|||||||
end
|
end
|
||||||
|
|
||||||
let λ = BN.λ
|
let λ = BN.λ
|
||||||
λ.(reshape(γ, affine_shape...) .* ((x .- μ) ./ sqrt.(σ² .+ BN.ϵ)) .+ reshape(β, affine_shape...))
|
# Break this up with a temporary variable to fix GPU launch bug
|
||||||
|
# https://github.com/FluxML/Flux.jl/issues/385
|
||||||
|
temp = reshape(γ, affine_shape...) .* ((x .- μ) ./ sqrt.(σ² .+ BN.ϵ))
|
||||||
|
return λ.(temp .+ reshape(β, affine_shape...))
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user