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

Update python/mlx/nn/layers/normalization.py

Browse Source 
Co-authored-by: Robert McCraith <mccraithrobert@gmail.com>
This commit is contained in:
__mo_san__ 2023-12-23 15:20:08 +01:00 committed by m0saan
parent 019a85511c
commit c68a472b83
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -234,8 +234,8 @@ class BatchNorm(Module):
self.bias = mx.zeros((num_features,))
if self.track_running_stats:
self.running_mean = mx.zeros((num_features,))
self.running_var = mx.ones((num_features,))
self._running_mean = mx.zeros((num_features,))
self._running_var = mx.ones((num_features,))
def _extra_repr(self):
return f"{self.num_features}, eps={self.eps}, momentum={self.momentum}, affine={'weight' in self}, track_running_stats={self.track_running_stats}"