Merge pull request #275 from alin-asapp/master

Fixed bug in scheduler logging
asappresearch · Mar 2, 2020 · 015bc93 · 015bc93
2 parents ab908dc + e2522de
commit 015bc93
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/flambe/learn/train.py b/flambe/learn/train.py
@@ -302,10 +302,13 @@ def _train_step(self) -> None:
                 # Optimize
                 self.optimizer.step()
 
+                # Log the learning rate
+                if self.iter_scheduler is not None or self.scheduler is not None:
+                    lr = self.optimizer.param_groups[0]['lr']  # type: ignore
+                    log(f'{log_prefix}/LR', lr, global_step)
+
                 # Update iter scheduler
                 if self.iter_scheduler is not None:
-                    learning_rate = self.iter_scheduler.get_lr()[0]  # type: ignore
-                    log(f'{log_prefix}/LR', learning_rate, global_step)
                     self.iter_scheduler.step()  # type: ignore
 
                 # Zero the gradients when exiting a train step