From 2f0374db54d437b16346ddadfa497d16bf6602f0 Mon Sep 17 00:00:00 2001 From: Juhan Bae Date: Tue, 19 Sep 2023 14:25:45 -0400 Subject: [PATCH] minor --- algorithmic_efficiency/workloads/criteo1tb/workload.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/algorithmic_efficiency/workloads/criteo1tb/workload.py b/algorithmic_efficiency/workloads/criteo1tb/workload.py index 2597e8904..65354f822 100644 --- a/algorithmic_efficiency/workloads/criteo1tb/workload.py +++ b/algorithmic_efficiency/workloads/criteo1tb/workload.py @@ -134,12 +134,18 @@ def _eval_model_on_split(self, global_batch_size=global_batch_size, repeat_final_dataset=False) loss = 0.0 + size = 0 for i in range(num_batches): eval_batch = next(self._eval_iters[split]) - if i == (num_batches - 1): - print(eval_batch.get('weights')) + # if i == (num_batches - 1): + # print(eval_batch.get('weights')) loss += self._eval_batch(params, eval_batch) + size += eval_batch.get('weights').sum() if USE_PYTORCH_DDP: dist.all_reduce(loss) + dist.all_reduce(size) + + print("SIZE") + print(size) mean_loss = loss.item() / num_examples return {'loss': mean_loss}