Use next instead of get_batch

NVIDIA · Mar 12, 2024 · 0021bb4 · 0021bb4
1 parent 55214c8
commit 0021bb4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -976,7 +976,7 @@ def get_forward_output_and_loss_func(self, validation_step=False, tuning=False):
         def fwd_output_and_loss_func(dataloader_iter, model, checkpoint_activations_all_layers=None):
 
             # Get data batch
-            batch = self.get_batch(dataloader_iter, tuning)
+            batch = next(dataloader_iter)
 
             # Transfer needed data to GPU
             required_keys = set()