register_hook

Author: rohitgr7

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -505,7 +505,8 @@ def _update_learning_rates(
                 self.scheduler_progress.increment_ready()
 
                 # update LR
-                self.trainer.lightning_module.lr_scheduler_step(
+                self.trainer._call_lightning_module_hook(
+                    "lr_scheduler_step",
                     lr_scheduler["scheduler"],
                     optimizer_idx=lr_scheduler["opt_idx"],
                     monitor_val=monitor_val,

pytorch_lightning/trainer/connectors/logger_connector/fx_validator.py

@@ -43,6 +43,7 @@ class _LogOptions(TypedDict):
         "optimizer_step": _LogOptions(
             allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
         ),
+        "lr_scheduler_step": None,
         "on_before_zero_grad": _LogOptions(
             allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
         ),

tests/trainer/logging_/test_logger_connector.py

@@ -233,6 +233,7 @@ def test_fx_validator_integration(tmpdir):
         "configure_callbacks": "You can't",
         "on_validation_model_eval": "You can't",
         "on_validation_model_train": "You can't",
+        "lr_scheduler_step": "You can't",
         "summarize": "not managed by the `Trainer",
     }
     model = HookedModel(not_supported)