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

Fix BN stats to not expand shape (#409)

Browse Source 
* fix BN stats to not expand shape

* nit
This commit is contained in:
Awni Hannun
2024-01-09 11:54:51 -08:00
committed by GitHub
parent 753867123d
commit e9ca65c939
2 changed files with 17 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -333,8 +333,8 @@ class BatchNorm(Module):
"""
reduction_axes = tuple(range(0, x.ndim - 1))
mean = mx.mean(x, axis=reduction_axes, keepdims=True)
var = mx.var(x, axis=reduction_axes, keepdims=True)
mean = mx.mean(x, axis=reduction_axes)
var = mx.var(x, axis=reduction_axes)
return mean, var