minor

mlcommons · Sep 19, 2023 · 2f0374d · 2f0374d
1 parent 5684336
commit 2f0374d
Showing 1 changed file with 8 additions and 2 deletions.
diff --git 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}