set init=0 in stats update

FluxML · Sep 30, 2023 · 2ce9c42 · 2ce9c42
1 parent 5b49a64
commit 2ce9c42
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/normalization.jl b/src/normalization.jl
@@ -138,7 +138,7 @@ function update_running_stats!(stats::RunningStats, x, μ, σ², reduce_dims::Di
     V = eltype(σ²)
     momentum = stats.momentum
     res_mtm = one(V) - momentum
-    m = prod(size(x, i) for i in reduce_dims)
+    m = prod(size(x, i) for i in reduce_dims; init = 1)
     correction = m / (m - one(V))
 
     running_mean, running_var = stats.mean, stats.variance