Pass mock to GPTDatasetConfig

Signed-off-by: Alexandros Koumparoulis <[email protected]>
NVIDIA · Sep 10, 2024 · 6e987de · 6e987de
1 parent 4259169
commit 6e987de
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
@@ -1598,6 +1598,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(