Edit num_training_batches

Signed-off-by: Abhishree <[email protected]>
NVIDIA · Sep 13, 2023 · 349bc13 · 349bc13
1 parent e269a7b
commit 349bc13
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py
@@ -1098,8 +1098,11 @@ def init_train_tqdm(self):
 
     def on_train_epoch_start(self, trainer, *_):
         if trainer.max_steps != -1:
-            num_training_batches = min(trainer.max_steps, trainer.num_training_batches)
+            #num_training_batches = min(trainer.max_steps, trainer.num_training_batches)
+            # update with max_steps for both training from scratch and resuming
+            num_training_batches = trainer.max_steps
         else:
+            # TODO: fix for resume with multiple epochs
             num_training_batches = trainer.num_training_batches
         self.train_progress_bar.reset(num_training_batches)
         self.train_progress_bar.initial = 0