Merge pull request #651 from RasmusOrsoe/display_current_learning_rate

Log current lr in `StandardModel`
graphnet-team · Feb 6, 2024 · 5c93a7a · 5c93a7a
2 parents d30bf71 + 4042375
commit 5c93a7a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/graphnet/models/standard_model.py b/src/graphnet/models/standard_model.py
@@ -273,6 +273,9 @@ def training_step(
             on_step=False,
             sync_dist=True,
         )
+
+        current_lr = self.trainer.optimizers[0].param_groups[0]["lr"]
+        self.log("lr", current_lr, prog_bar=True, on_step=True)
         return loss
 
     def validation_step(