Hardcodings for Grain Pipeline

AI-Hypercomputer · Aug 20, 2024 · 90d7304 · 90d7304
1 parent bf33d7c
commit 90d7304
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/MaxText/input_pipeline/input_pipeline_interface.py b/MaxText/input_pipeline/input_pipeline_interface.py
@@ -129,7 +129,8 @@ def make_mixed_train_iterator(config, mesh, add_bos, add_eos):
     if config.dataset_type == "tfds":
       return make_tfds_iterator(config, mesh, add_bos, add_eos, process_indices)
     elif config.dataset_type == "grain":
-      return make_grain_iterator(config, mesh, add_bos, add_eos, process_indices)
+      # Hardcoding this for now
+      return make_grain_iterator(config, mesh, False, False, process_indices)
     elif config.dataset_type == "hf":
       return make_hf_iterator(config, mesh, add_bos, add_eos, process_indices)
   else: