Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/loops/optimizer/optimizer_loop.py (Python)

Code before:
        self.optim_progress: OptimizationProgress = OptimizationProgress()

        self._skip_backward: bool = False
        self._batch_idx: int = 0
        self._optimizers: List[Optimizer] = []
        self._hiddens: Optional[Any] = None

    @property
    def done(self) -> bool:
        """Returns ``True`` when the last optimizer in the sequence has run."""
        return self.optim_progress.optimizer_idx >= len(self._optimizers)

    def connect(self, **kwargs: "Loop") -> None:
        raise NotImplementedError(f"{self.__class__.__name__} does not connect any child loops.")

    def reset(self) -> None:
        if not self.restarting:
            self.optim_progress.optimizer_idx = 0
        self.outputs = [[] for _ in range(len(self.trainer.optimizers))]

    def on_run_start(self, batch: Any, optimizers: List[Optimizer], batch_idx: int) -> None:  # type: ignore[override]
        self._batch_idx = batch_idx
        self._optimizers = optimizers

    def advance(self, batch: Any, *args, **kwargs) -> None:  # type: ignore[override]
        result = self._run_optimization(
            batch,
            self._batch_idx,
            self._optimizers[self.optim_progress.optimizer_idx],
            self.optim_progress.optimizer_idx,
        )
        if result:
            self.outputs[self.optim_progress.optimizer_idx].append(deepcopy(result.result_collection))

        self.optim_progress.optimizer_idx += 1

    def on_run_end(self) -> _OUTPUTS_TYPE:
        outputs, self.outputs = self.outputs, []  # free memory
        return outputs

    def backward(
        self,
        loss: Tensor,
        optimizer: torch.optim.Optimizer,
        opt_idx: int,
        *args: Any,
        **kwargs: Any,
    ) -> Tensor:
        """Performs the backward step.

        Args:

Review comment (suggestion):
@awaelchli `_run_optimization` could have a `hiddens` argument which gets passed all the way down to `_training_step`.

Or `_training_step` could access `self._hiddens` directly. What do you prefer?

Is there a reason why you were setting `self._hiddens` here and accessing it again in `_run_optimization`?

context:
In the future, `advance` will be the one to set `self._hiddens = result.hiddens` after `_run_optimization`

Code after:
        self.optim_progress: OptimizationProgress = OptimizationProgress()

        self._skip_backward: bool = False
        self._batch_idx: int = 0
        self._optimizers: List[Optimizer] = []
        self._hiddens: Optional[Any] = None

    @property
    def done(self) -> bool:
        """Returns ``True`` when the last optimizer in the sequence has run."""
        return self.optim_progress.optimizer_idx >= len(self._optimizers)

    def connect(self, **kwargs: "Loop") -> None:
        raise NotImplementedError(f"{self.__class__.__name__} does not connect any child loops.")

    def reset(self) -> None:
        if not self.restarting:
            self.optim_progress.optimizer_idx = 0
        self.outputs = [[] for _ in range(len(self.trainer.optimizers))]

    def on_run_start(self, batch: Any, optimizers: List[Optimizer], batch_idx: int) -> None:  # type: ignore[override]
        self._batch_idx = batch_idx
        self._optimizers = optimizers

    def advance(self, batch: Any, *args, **kwargs) -> None:  # type: ignore[override]
        result = self._run_optimization(
            batch,
            self._batch_idx,
            self._optimizers[self.optim_progress.optimizer_idx],
            self.optim_progress.optimizer_idx,
        )
        if result.result_collection is not None:
            self.outputs[self.optim_progress.optimizer_idx].append(deepcopy(result.result_collection))

        self.optim_progress.optimizer_idx += 1

    def on_run_end(self) -> _OUTPUTS_TYPE:
        outputs, self.outputs = self.outputs, []  # free memory
        return outputs

    def backward(
        self, loss: Tensor, optimizer: torch.optim.Optimizer, opt_idx: int, *args: Any, **kwargs: Any
    ) -> Tensor:
        """Performs the backward step.

        Args:
            loss: The loss value to back-propagate on
            optimizer: Current optimizer being used
            opt_idx: Index of the current optimizer being used
        """
        self.trainer.accelerator.backward(loss, optimizer, opt_idx, *args, **kwargs)