Address the merge conflict properly

Signed-off-by: John St John <jstjohn@nvidia.com>
NVIDIA · Aug 15, 2024 · 6f4e2d9 · 6f4e2d9
1 parent 18e9469
commit 6f4e2d9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/nemo/lightning/pytorch/plugins/data_sampler.py b/nemo/lightning/pytorch/plugins/data_sampler.py
@@ -96,8 +96,7 @@ def on_megatron_step_end(self, trainer: pl.Trainer, pl_module: pl.LightningModul
 
         self.prev_global_batch_size = self.current_global_batch_size
 
-        # TODO: Add consumed samples
-        consumed_samples = self.compute_consumed_samples(trainer.global_step + 1 - self.init_consumed_samples)
+        consumed_samples = self.compute_consumed_samples(trainer.global_step + 1 - self.init_global_step)
         if self.output_log:
             # You may need to turn off logging, for example when doing trainer.predict(model, data)
             pl_module.log(