def _save_grad_output(self, mod, _, grad_output):
        """Saves grad on output of layer to grad is scaled with batch_size since gradient is spread over samples in
        mini batch."""
        batch_size = grad_output[0].shape[0]
        if mod.training:
            self.state[mod]["grad"] = grad_output[0] * batch_size
Saves grad on output of layer to grad is scaled with batch_size since gradient is spread over samples in
        mini batch.
tests/tests_pytorch/core/test_lightning_optimizer.py