fix .get_last_lr() error refer to ildoonet/pytorch-gradual-warmup-lr#14

Author: vvv

warmup_scheduler/run.py

@@ -10,7 +10,7 @@
     optim = SGD(model, 0.1)
 
     # scheduler_warmup is chained with schduler_steplr
-    scheduler_steplr = StepLR(optim, step_size=10, gamma=0.1)
+    scheduler_steplr = StepLR(optim, step_size=5, gamma=0.1)
     scheduler_warmup = GradualWarmupScheduler(optim, multiplier=1, total_epoch=5, after_scheduler=scheduler_steplr)
 
     # this zero gradient update is needed to avoid a warning message, issue #8.

warmup_scheduler/scheduler.py

@@ -28,7 +28,7 @@ def get_lr(self):
                 if not self.finished:
                     self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]
                     self.finished = True
-                return self.after_scheduler.get_last_lr()
+                return [group['lr'] for group in self.optimizer.param_groups]
             return [base_lr * self.multiplier for base_lr in self.base_lrs]
 
         if self.multiplier == 1.0:
@@ -57,7 +57,7 @@ def step(self, epoch=None, metrics=None):
                     self.after_scheduler.step(None)
                 else:
                     self.after_scheduler.step(epoch - self.total_epoch)
-                self._last_lr = self.after_scheduler.get_last_lr()
+                self._last_lr = [group['lr'] for group in self.optimizer.param_groups]
             else:
                 return super(GradualWarmupScheduler, self).step(epoch)
         else: