minor

mlcommons · Sep 19, 2023 · 1d82836 · 1d82836
1 parent 2f0374d
commit 1d82836
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/algorithmic_efficiency/workloads/criteo1tb/workload.py b/algorithmic_efficiency/workloads/criteo1tb/workload.py
@@ -135,8 +135,7 @@ def _eval_model_on_split(self,
           repeat_final_dataset=False)
     loss = 0.0
     size = 0
-    for i in range(num_batches):
-      eval_batch = next(self._eval_iters[split])
+    for eval_batch in self._eval_iters[split]:
       # if i == (num_batches - 1):
       #   print(eval_batch.get('weights'))
       loss += self._eval_batch(params, eval_batch)