min(n_proc , ...)

microsoft · Nov 12, 2024 · ec885eb · ec885eb
1 parent bc14f13
commit ec885eb
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/mttl/datamodule/mt_seq_to_seq_module.py b/mttl/datamodule/mt_seq_to_seq_module.py
@@ -138,7 +138,9 @@ def apply_source_template(dataset, source_template):
 class FlatMultiTaskModule(DataModule):
     def setup_dataset(self):
         self.dataset = DatasetLibrary.pull_dataset_with_retry(self.config.dataset)
-        n_proc = int(os.environ.get("MTTL_NUM_PROC_DATASETS", 16))
+        n_proc = min(
+            len(self.dataset), int(os.environ.get("MTTL_NUM_PROC_DATASETS", 16))
+        )
 
         if "split" not in self.dataset.column_names["train"]:
             logger.warning(