code: def __init__(self, model):
      self._fwd_handles = []
      self._bwd_handles = []
      self.params = []
      for _, mod in model.named_modules():
          mod_class = mod.__class__.__name__
          if mod_class != "Linear":
              continue
  
          handle = mod.register_forward_pre_hook(
              self._save_input
          )  # save the inputs
          self._fwd_handles.append(
              handle
          )  # collect forward-save-input hooks in list
          handle = mod.register_backward_hook(
              self._save_grad_output
          )  # save the gradients
          self._bwd_handles.append(
              handle
          )  # collect backward-save-grad hook in list
  
          # save the parameters
          params = [mod.weight]
          if mod.bias is not None:
              params.append(mod.bias)
  
          # save a param_group for each module
          d = {"params": params, "mod": mod, "layer_type": mod_class}
          self.params.append(d)
  
      super().__init__(self.params, {"lr": 0.01})
  
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 1 to create a new OptimizerWithHooks object from inherited base classes: Optimizer with example: obj = OptimizerWithHooks(model)