integrate #7563

Author: awaelchli

pytorch_lightning/loops/batch_loop.py

@@ -124,7 +124,6 @@ def _run_optimization(self, batch_idx, split_idx, split_batch, opt_idx=0, optimi
             # -------------------
             # calculate loss (train step + train step end)
             # -------------------
-
             # automatic_optimization=True: perform ddp sync only when performing optimizer_step
             # automatic_optimization=False: don't block synchronization here
             with self.block_ddp_sync_behaviour():
@@ -137,6 +136,9 @@ def _run_optimization(self, batch_idx, split_idx, split_batch, opt_idx=0, optimi
         else:
             if self.trainer.lightning_module.automatic_optimization:
                 self.optimizer_step(optimizer, opt_idx, batch_idx, closure)
+                if len(self.trainer.optimizers) > 1:
+                    # revert back to previous state
+                    self.trainer.lightning_module.untoggle_optimizer(opt_idx)
             else:
                 result = self.training_step(split_batch, batch_idx, opt_idx, self._hiddens)
 
@@ -448,10 +450,6 @@ def training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer,
                         "training_step returned None. If this was on purpose, ignore this warning..."
                     )
 
-                if len(self.trainer.optimizers) > 1:
-                    # revert back to previous state
-                    self.trainer.lightning_module.untoggle_optimizer(opt_idx)
-
         return result
 
     def _check_finite(self, loss: torch.Tensor) -> None: