zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 18:26:37 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix average_stats

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-02-28 16:00:06 -08:00
parent d20413a54d
commit 66b630b4f6
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
cifar/main.py
View File

@ -54,10 +54,7 @@ def train_epoch(model, train_iter, optimizer, epoch):
with mx.stream(mx.cpu): with mx.stream(mx.cpu):
stats = mx.distributed.all_sum(mx.array(stats)) stats = mx.distributed.all_sum(mx.array(stats))
count = mx.distributed.all_sum(count) count = mx.distributed.all_sum(count)
mx.eval(stats, count) return (stats / count).tolist()
count = count.item()
return [s / count for s in stats.tolist()]
state = [model.state, optimizer.state] state = [model.state, optimizer.state]