docstr: Saves grad on output of layer to grad is scaled with batch_size since gradient is spread over samples in
  mini batch.
docstr_tok: ["Saves","grad","on","output","of","layer","to","grad","is","scaled","with","batch_size","since","gradient","is","spread","over","samples","in","mini","batch","."]
code: 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
  
imports: from copy import deepcopy
  from unittest.mock import DEFAULT, Mock, patch
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.core.optimizer import LightningOptimizer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from lightning.pytorch.loops.optimization.automatic import Closure
  from lightning.pytorch.tuner.tuning import Tuner
  from torch.optim import SGD, Adam, Optimizer
label_desc: Use image node_id 3 for calling the OptimizerWithHooks obj's underlying member method code with example usage: obj._save_grad_output(mod, _, grad_output) without return types