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

Fix the throughput calculation

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-02-25 17:20:15 -08:00
parent 8a76b421a0
commit 14faec4ca2
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
cifar/main.py
View File

@ -81,7 +81,10 @@ def train_epoch(model, train_iter, optimizer, epoch):
samples_per_sec += x.shape[0] / (toc - tic)
count += 1
if batch_counter % 10 == 0:
l, a, s = average_stats([losses, accuracies, samples_per_sec], count)
l, a, s = average_stats(
[losses, accuracies, world.size() * samples_per_sec],
count,
)
print_zero(
world,
" | ".join(
@ -94,7 +97,7 @@ def train_epoch(model, train_iter, optimizer, epoch):
),
)
return average_stats([losses, accuracies, samples_per_sec], count)
return average_stats([losses, accuracies, world.size() * samples_per_sec], count)
def test_epoch(model, test_iter, epoch):