Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/training_loop.py (Python)

Code before:
                optimizers = [optimizers[0]]

            # loop over optimizers
            for opt_idx, optimizer in optimizers:
                # make sure only the gradients of the current optimizer's parameters are calculated
                # in the training step to prevent dangling gradients in multiple-optimizer setup.
                if self.automatic_optimization and len(self.trainer.optimizers) > 1:
                    model = self.trainer.get_model()
                    model.toggle_optimizer(optimizer, opt_idx)

                if not (accumulation_done or is_final_batch):
                    # For gradient accumulation

                    # -------------------
                    # calculate loss (train step + train step end)
                    # -------------------

                    # no ddp sync at the beginning of forward or backward due to parameter changes in this or last step required
                    no_sync = self._updated_model_last_step and isinstance(self.trainer.model, torch.nn.parallel.DistributedDataParallel)
                    if no_sync:
                        self.trainer.model.no_sync.__enter__()

                    self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)

                    if no_sync:
                        self.trainer.model.__exit__()

                    batch_outputs = self._process_closure_result(
                        batch_callback_metrics=batch_callback_metrics,
                        batch_log_metrics=batch_log_metrics,
                        batch_outputs=batch_outputs,
                        opt_idx=opt_idx,
                    )

                    self._updated_model_last_step = False

                # ------------------------------
                # BACKWARD PASS
                # ------------------------------
                # gradient update with accumulated gradients

                else:
                    self._updated_model_last_step = True

                    if self.automatic_optimization:

                        def train_step_and_backward_closure():
                            result = self.training_step_and_backward(
                                split_batch,
                                batch_idx,
                                opt_idx,

Review comment (suggestion):
```suggestion
                        self.trainer.model.no_sync.__exit__()
```

Code after:
            if not self.automatic_optimization:
                optimizers = [optimizers[0]]

            # loop over optimizers
            for opt_idx, optimizer in optimizers:
                # make sure only the gradients of the current optimizer's parameters are calculated
                # in the training step to prevent dangling gradients in multiple-optimizer setup.
                if self.automatic_optimization and len(self.trainer.optimizers) > 1:
                    model = self.trainer.get_model()
                    model.toggle_optimizer(optimizer, opt_idx)

                if should_accumulate:
                    # For gradient accumulation

                    # -------------------
                    # calculate loss (train step + train step end)
                    # -------------------

                    # perform dpp sync only when performing optimizer_step
                    with self.block_ddp_sync_behaviour():
                        self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)

                    batch_outputs = self._process_closure_result(
                        batch_callback_metrics=batch_callback_metrics,
                        batch_log_metrics=batch_log_metrics,
                        batch_outputs=batch_outputs,
                        opt_idx=opt_idx,
                    )

                # ------------------------------
                # BACKWARD PASS
                # ------------------------------
                # gradient update with accumulated gradients

                else:
                    if self.automatic_optimization:

                        def train_step_and_backward_closure():
                            result = self.training_step_and_backward(
                                split_batch,
                                batch_idx,
                                opt_idx,
                                optimizer,
                                self.trainer.hiddens
                            )
                            return None if result is None else result.loss

                        # optimizer step
                        self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)

                    else: