add toggled_optimizer to LightningModule

Author: rustamzh

src/lightning/pytorch/core/module.py

@@ -1141,6 +1141,30 @@ def untoggle_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) ->
         # save memory
         self._param_requires_grad_state = {}
 
+    @contextmanager
+    def toggled_optimizer(self, optimizer: Union[Optimizer, LightningOptimizer]) -> Generator:
+        """Makes sure only the gradients of the current optimizer's parameters are calculated in the training step to
+        prevent dangling gradients in multiple-optimizer setup. Combines :meth:`toggle_optimizer` and 
+        :meth:`untoggle_optimizer` into context manager.
+
+        Args:
+            optimizer: The optimizer to untoggle.
+
+        Example::
+
+            def training_step(...):
+                opt = self.optimizers()
+                with self.toggled_optimizer(opt):
+                    loss = ...
+                    opt.zero_grad()
+                    self.manual_backward(loss)
+                    opt.step()
+        """
+        try:
+            yield self.toggle_optimizer(optimizer)
+        finally:
+            self.untoggle_optimizer(optimizer)
+
     def clip_gradients(
         self,
         optimizer: Optimizer,

tests/tests_pytorch/core/test_lightning_module.py

@@ -118,6 +118,21 @@ def test_1_optimizer_toggle_model():
     model.untoggle_optimizer(optimizer)
     assert not model._param_requires_grad_state
 
+def test_1_optimizer_toggle_model_context_manager():
+    """Test toggle_model runs when only one optimizer is used."""
+    model = BoringModel()
+    trainer = Mock()
+    model.trainer = trainer
+    params = model.parameters()
+    optimizer = torch.optim.SGD(params, lr=0.1)
+    trainer.optimizers = [optimizer]
+
+    assert not model._param_requires_grad_state
+    # toggle optimizer was failing with a single optimizer
+    with model.toggled_optimizer(optimizer):
+        assert model._param_requires_grad_state
+    assert not model._param_requires_grad_state
+
 
 def test_toggle_untoggle_2_optimizers_no_shared_parameters(tmp_path):
     class TestModel(BoringModel):