add dataset_root (#11837)

NVIDIA · Jan 17, 2025 · ad807ae · ad807ae
1 parent ca4e4f0
commit ad807ae
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/collections/llm/bert/data/specter.py b/nemo/collections/llm/bert/data/specter.py
@@ -43,6 +43,7 @@ class SpecterDataModule(FineTuningDataModule, IOMixin):
 
     def __init__(
         self,
+        dataset_root: str = None,
         seq_length: int = 512,
         tokenizer: Optional["TokenizerSpec"] = None,
         micro_batch_size: int = 4,
@@ -61,7 +62,7 @@ def __init__(
         self.delete_raw = delete_raw
 
         super().__init__(
-            dataset_root=get_dataset_root("specter"),
+            dataset_root=get_dataset_root("specter") if dataset_root is None else dataset_root,
             seq_length=seq_length,
             tokenizer=tokenizer,
             micro_batch_size=micro_batch_size,