From 14faec4ca2794d56793a7080b8f3c5fe073c087c Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 25 Feb 2025 17:20:15 -0800 Subject: [PATCH] Fix the throughput calculation --- cifar/main.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cifar/main.py b/cifar/main.py index 27074133..7eb6efdf 100644 --- a/cifar/main.py +++ b/cifar/main.py @@ -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):