Fix and CHANGELOG

Author: carmocca

CHANGELOG.md

@@ -77,6 +77,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `opt_idx` to scheduler config if not assigned by user ([#11247](https://github.com/PyTorchLightning/pytorch-lightning/pull/11247))
 
+- Added support for optimizer step progress tracking with manual optimization ([#11848](https://github.com/PyTorchLightning/pytorch-lightning/pull/11848))
+
+
 
 - Return the output of the `optimizer.step`. This can be useful for `LightningLite` users, manual optimization users, or users overriding `LightningModule.optimizer_step` ([#11711](https://github.com/PyTorchLightning/pytorch-lightning/pull/11711))
 

pytorch_lightning/loops/optimization/manual_loop.py

@@ -16,6 +16,7 @@
 
 from torch import Tensor
 
+from pytorch_lightning.core.optimizer import do_nothing_closure
 from pytorch_lightning.loops import Loop
 from pytorch_lightning.loops.optimization.closure import OutputResult
 from pytorch_lightning.loops.utilities import _build_training_step_kwargs, _extract_hiddens
@@ -76,7 +77,7 @@ class ManualOptimization(Loop[_OUTPUTS_TYPE]):
     def __init__(self) -> None:
         super().__init__()
         # since manual optimization does not track lr scheduler or optimizer frequencies, we use a simpler progress than
-        # `OptimizerProgress`
+        # `OptimizationProgress`
         self.optim_step_progress = Progress.from_defaults(ReadyCompletedTracker)
 
         self._done: bool = False
@@ -137,6 +138,10 @@ def advance(self, batch: Any, batch_idx: int) -> None:  # type: ignore[override]
     def on_run_end(self) -> _OUTPUTS_TYPE:
         """Returns the result of this loop, i.e., the post-processed outputs from the training step."""
         output, self._output = self._output, {}  # free memory
+        # reset logic around the optimizer step
+        for i, lightning_optimizer in self.trainer.strategy._lightning_optimizers.items():
+            lightning_optimizer.on_before_step = do_nothing_closure
+            lightning_optimizer.on_after_step = do_nothing_closure
         return output
 
     def _on_before_step(self) -> None: