add epoch to lr_scheduler.step() call

fix issue victoresque#68 of victoresque/pytorch-template
SunQpark · Jun 12, 2020 · ce38fb6 · ce38fb6
1 parent d90f283
commit ce38fb6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/srcs/trainer/trainer.py b/srcs/trainer/trainer.py
@@ -70,7 +70,7 @@ def _train_epoch(self, epoch):
             log.update(**{'val_'+k : v for k, v in val_log.items()})
 
         if self.lr_scheduler is not None:
-            self.lr_scheduler.step()
+            self.lr_scheduler.step(epoch)
         return log
 
     def _valid_epoch(self, epoch):