add docs for toggled_optimizer to LightningModule

Author: rustamzh

docs/source-pytorch/conf.py

@@ -487,6 +487,7 @@ def _load_py_module(name: str, location: str) -> ModuleType:
     ("py:meth", "setup"),
     ("py:meth", "test_step"),
     ("py:meth", "toggle_optimizer"),
+    ("py:meth", "toggled_optimizer"),
     ("py:class", "torch.ScriptModule"),
     ("py:class", "torch.distributed.fsdp.fully_sharded_data_parallel.CPUOffload"),
     ("py:class", "torch.distributed.fsdp.fully_sharded_data_parallel.MixedPrecision"),

docs/source-pytorch/model/manual_optimization.rst

@@ -17,7 +17,7 @@ To manually optimize, do the following:
   * ``optimizer.zero_grad()`` to clear the gradients from the previous training step
   * ``self.manual_backward(loss)`` instead of ``loss.backward()``
   * ``optimizer.step()`` to update your model parameters
-  * ``self.toggle_optimizer()`` and ``self.untoggle_optimizer()`` if needed
+  * ``self.toggle_optimizer()`` and ``self.untoggle_optimizer()``, or ``self.toggled_optimizer()`` if needed
 
 Here is a minimal example of manual optimization.
 

src/lightning/pytorch/CHANGELOG.md

@@ -11,6 +11,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Add enable_autolog_hparams argument to Trainer ([#20593](https://github.com/Lightning-AI/pytorch-lightning/pull/20593))
 
+- Add `toggled_optimizer(optimizer)` method to the LightningModule, which is a context manager version of `toggle_optimize` and `untoggle_optimizer` 
+
 
 ### Changed
 

src/lightning/pytorch/core/module.py

@@ -1148,7 +1148,7 @@ def toggled_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) ->
         :meth:`untoggle_optimizer` into context manager.
 
         Args:
-            optimizer: The optimizer to untoggle.
+            optimizer: The optimizer to toggle.
 
         Example::