update callback to use new scheduler methods

Author: Chris-hughes10

pytorch_accelerated/callbacks.py

@@ -834,7 +834,7 @@ def on_training_run_start(self, trainer, **kwargs):
     def on_train_step_end(self, trainer, step: int, **kwargs):
         """Handle checkpoint saving and progress logging"""
 
-        total_steps = trainer.scheduler.get_current_step()
+        total_steps = trainer.scheduler.get_current_step() + 1
 
         # Skip if we've already saved at this step
         if total_steps == self.last_checkpoint_step:
@@ -861,7 +861,7 @@ def on_train_step_end(self, trainer, step: int, **kwargs):
     def on_training_run_end(self, trainer, **kwargs):
         """Save final checkpoint if we haven't already"""
         # Get the final step number
-        total_steps = trainer.run_config.max_num_train_steps
+        total_steps = trainer.scheduler.get_current_step() + 1
 
         # If we haven't saved the final checkpoint yet
         if self.last_checkpoint_step != total_steps: