Pass mock to GPTDatasetConfig (NVIDIA#10435)

Signed-off-by: Alexandros Koumparoulis <[email protected]> Co-authored-by: Pablo Garay <[email protected]> Signed-off-by: Hainan Xu <[email protected]>
hainan-xv · Nov 5, 2024 · f3c632e · f3c632e
1 parent 27a3dc5
commit f3c632e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -1622,6 +1622,7 @@ def build_train_valid_test_datasets(self):
                 dataset_type = GPTFIMDataset
             else:
                 dataset_config = GPTDatasetConfig(**kwargs)
+                dataset_config.mock = mock_dataset
                 dataset_type = MockGPTDataset if mock_dataset else GPTDataset
 
             self._train_ds, self._validation_ds, self._test_ds = BlendedMegatronDatasetBuilder(