def optimizer_step(
        self,
        epoch: int,
        batch_idx: int,
        optimizer: Union[Optimizer, LightningOptimizer],
        optimizer_closure: Optional[Callable[[], Any]] = None,
    ) -> None:
        r"""Override this method to adjust the default way the :class:`~lightning.pytorch.trainer.trainer.Trainer` calls
        the optimizer.

        By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example.
        This method (and ``zero_grad()``) won't be called during the accumulation phase when
        ``Trainer(accumulate_grad_batches != 1)``. Overriding this hook has no benefit with manual optimization.

        Args:
            epoch: Current epoch
            batch_idx: Index of current batch
            optimizer: A PyTorch optimizer
            optimizer_closure: The optimizer closure. This closure must be executed as it includes the
                calls to ``training_step()``, ``optimizer.zero_grad()``, and ``backward()``.

        Examples::

            def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
                # Add your custom logic to run directly before `optimizer.step()`

                optimizer.step(closure=optimizer_closure)

                # Add your custom logic to run directly after `optimizer.step()`

        """
        optimizer.step(closure=optimizer_closure)
Override this method to adjust the default way the :class:`~lightning.pytorch.trainer.trainer.Trainer` calls
        the optimizer.

        By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example.
        This method (and ``zero_grad()``) won't be called during the accumulation phase when
        ``Trainer(accumulate_grad_batches != 1)``. Overriding this hook has no benefit with manual optimization.

        Args:
            epoch: Current epoch
            batch_idx: Index of current batch
            optimizer: A PyTorch optimizer
            optimizer_closure: The optimizer closure. This closure must be executed as it includes the
                calls to ``training_step()``, ``optimizer.zero_grad()``, and ``backward()``.

        Examples::

            def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
                # Add your custom logic to run directly before `optimizer.step()`

                optimizer.step(closure=optimizer_closure)

                # Add your custom logic to run directly after `optimizer.step()`

        
src/lightning/pytorch/core/module.py