Re-fix type stability

LuxDL · Sep 2, 2023 · 1158df8 · 1158df8
1 parent b424886
commit 1158df8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.3.2"
+version = "0.3.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

diff --git a/src/impl/normalization.jl b/src/impl/normalization.jl
@@ -3,12 +3,13 @@ function _update_normalization_statistics(x::AA{<:Real, N}, running_mean::AA{<:R
     running_var::AA{<:Real, N}, batchmean::AA{<:Real, N}, batchvar::AA{<:Real, N},
     momentum::Real, ::Val{reduce_dims}) where {N, reduce_dims}
     m = eltype(x)(prod(Base.Fix1(size, x), reduce_dims))
+    m_ = m / (m - one(m))
     if last(reduce_dims) != N
         batchmean = mean(batchmean; dims=N)
         batchvar = mean(batchvar; dims=N)
     end
     running_mean = @. (1 - momentum) * running_mean + momentum * batchmean
-    running_var = @. (1 - momentum) * running_var + momentum * batchvar * (m / (m - one(m)))
+    running_var = @. (1 - momentum) * running_var + momentum * batchvar * m_
     return (running_mean, running_var)
 end