autounit - extract lr scheduler step to a separate method (#778)

Summary: Pull Request resolved: #778

Reviewed By: JKSenthil

Differential Revision: D55754227

fbshipit-source-id: 4d2ea236359001585debcd954b1a11606cafae42

Author: galrotem

torchtnt/framework/auto_unit.py

@@ -408,6 +408,8 @@ class AutoUnit(
     - ``on_eval_step_end``
     - ``on_predict_step_end``
 
+    The user can also override the LR step method, ``step_lr_scheduler``, in case they want to have custom logic.
+
     Then use with the :py:func:`~torchtnt.framework.train`, :py:func:`~torchtnt.framework.evaluate`, :py:func:`~torchtnt.framework.fit`, or
     :py:func:`~torchtnt.framework.predict` entry point as normal.
 
@@ -731,6 +733,12 @@ def on_predict_step_end(
         """
         pass
 
+    def step_lr_scheduler(self) -> None:
+        """
+        LR step method extracted to a method in case the user wants to override
+        """
+        none_throws(self.lr_scheduler).step()
+
     def _update_weights(self, state: State) -> Optional[torch.Tensor]:
         """
         Updates weights of the module, handles clip gradient norm, etc.
@@ -865,12 +873,11 @@ def _update_lr_and_swa(self, state: State, number_of_steps_or_epochs: int) -> No
                 none_throws(self.swa_scheduler).step()
         else:
             # optionally step lr scheduler if SWA not in use
-            lr_scheduler = self.lr_scheduler
-            if lr_scheduler:
+            if self.lr_scheduler:
                 with get_timing_context(
                     state, f"{self.__class__.__name__}.lr_scheduler_step"
                 ):
-                    lr_scheduler.step()
+                    self.step_lr_scheduler()
 
 
 def _validate_torch_compile_available() -> None: