Add num_batch configs

mlcommons · Sep 25, 2023 · 86ad0af · 86ad0af
1 parent cc8b820
commit 86ad0af
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/algorithmic_efficiency/workloads/criteo1tb/criteo1tb_pytorch/workload.py b/algorithmic_efficiency/workloads/criteo1tb/criteo1tb_pytorch/workload.py
@@ -136,8 +136,6 @@ def _build_input_queue(
       cache: Optional[bool] = None,
       repeat_final_dataset: Optional[bool] = None,
       num_batches: Optional[int] = None) -> Iterator[Dict[str, spec.Tensor]]:
-    del num_batches
-
     not_train = split != 'train'
     per_device_batch_size = int(global_batch_size / N_GPUS)
 
@@ -149,6 +147,7 @@ def _build_input_queue(
           split=split,
           data_dir=data_dir,
           global_batch_size=global_batch_size,
+          num_batches=num_batches,
           repeat_final_dataset=repeat_final_dataset)
     weights = None
     while True:

diff --git a/algorithmic_efficiency/workloads/criteo1tb/workload.py b/algorithmic_efficiency/workloads/criteo1tb/workload.py
@@ -95,13 +95,13 @@ def _build_input_queue(
       repeat_final_dataset: Optional[bool] = None,
       num_batches: Optional[int] = None) -> Iterator[Dict[str, spec.Tensor]]:
     del cache
-    del num_batches
     ds = input_pipeline.get_criteo1tb_dataset(
         split=split,
         shuffle_rng=data_rng,
         data_dir=data_dir,
         num_dense_features=self.num_dense_features,
         global_batch_size=global_batch_size,
+        num_batches=num_batches,
         repeat_final_dataset=repeat_final_dataset)
 
     for batch in iter(ds):
@@ -132,6 +132,7 @@ def _eval_model_on_split(self,
           split=split,
           data_dir=data_dir,
           global_batch_size=global_batch_size,
+          num_batches=num_batches,
           repeat_final_dataset=True)
     loss = 0.0
     for _ in range(num_batches):