Hi All,
I was wondering if it were possible to save the grad attributes of my model for all my parameters? I’m currently using a custom optimizer in my problem and it requires using the .grad attribute (as I’m preconditioning gradients).
The way I save my model is via,
    torch.save({'epoch':epoch,
                'model_state_dict':net.state_dict(),
                'optim_state_dict':optim.state_dict(),
                'loss':loss}, model_path)

and is subsequently loaded via,
state_dict = torch.load(f=model_path, map_location=torch.device('cpu'))
net.load_state_dict(state_dict['model_state_dict']) #returns <All keys matched successfully>
optim.load_state_dict(state_dict['optim_state_dict']) #return nothing (I assume ok?)

If I load the model and print the grad attribute it doesn’t exist and hence causes my code to crash. So, I was wondering how exactly can I save these values? I printed the grad attribute via,
for name, param in net.named_parameters():
  print(param.grad)

The strange thing is I called loss.backward() before calling optim.step() so I thought that the grad attributes would be filled with a grad attribute? This is done in the standard way like,
optim.zero_grad()
loss_mean.backward()
optim.step()

Any help would be appreicated! Thank you!