zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-21 20:46:46 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

calc running mean and var only when training

Browse Source
This commit is contained in:
m0saan 2023-12-19 06:29:52 +01:00
parent c3c2fcf41d
commit d4bf9a2976
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/mlx/nn/layers/normalization.py
View File

@ -258,7 +258,7 @@ class BatchNorm1d(Module):
means = mx.mean(x, axis=(0, 2), keepdims=True) means = mx.mean(x, axis=(0, 2), keepdims=True)
var = mx.var(x, axis=(0, 2), keepdims=True) var = mx.var(x, axis=(0, 2), keepdims=True)
if self.track_running_stats: if self.track_running_stats and self.training:
self.running_mean = ( self.running_mean = (
1 - self.momentum 1 - self.momentum
) * self.running_mean + self.momentum * means.squeeze() ) * self.running_mean + self.momentum * means.squeeze()