Merge pull request #150 from TeamEpochGithub/149-bug-step-in-schedule…

…r-is-delayed-by-1 Fix scheduler step delay in torch trainer
TeamEpochGithub · May 16, 2024 · 984a786 · 984a786
2 parents 221eeb4 + a33ec46
commit 984a786
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/epochalyst/pipeline/model/training/torch_trainer.py b/epochalyst/pipeline/model/training/torch_trainer.py
@@ -625,7 +625,7 @@ def _train_one_epoch(
 
         # Step the scheduler
         if self.initialized_scheduler is not None:
-            self.initialized_scheduler.step(epoch=epoch)
+            self.initialized_scheduler.step(epoch=epoch + 1)
 
         # Remove the cuda cache
         torch.cuda.empty_cache()