Ok, so I’m making a custom optimizer based on KFAC and I noticed that if I resume my model from a checkpoint it fails because the grad attributes don’t exist. When I first initialise the optimizer, I append each module to a list and apply a forward and backward hook to each module. I also store the parameters of modules into the params list. This is done via,
  def _prepare_model(self):
    print("Adding parameters to module, marking nn.Linear layers!")
    for module in self.net.modules():
      if(module.__class__.__name__ in self._kfac_accepted_modules):
        print("module: ",module)
        self.modules.append(module)
        
        handle = module.register_forward_pre_hook(self._save_input)
        self._fwd_handles.append(handle)
        handle = module.register_full_backward_hook(self._save_output)
        self._bwk_handles.append(handle)
        
        params = [module.weight]
        if(module.bias is not None):
          params.append(module.bias)
        d = {'params':params, 'module':module}
        self.params.append(d)

Then when I go to precondition the gradient attribute I grab the shape of it via group['module'].weight.grad.shape. This didn’t work, but what did work was to change group['module'].weight.grad to group['params'][0].grad which does exist and gets created when loss.backward() gets called. This is done via the following, and is called on each group within optim.param_groups
def _precondition(self, group):
  module = group['module']

  g = group['params'][0].grad
  s = g.shape
  if(module.bias is not None):
    #gb = module.bias.grad       #doesn't work
    gb = group['params'][1].grad #does work
    g = torch.cat([g, gb.unsqueeze(1)], dim=1)

Does this clarify the workflow? Thank you!