Hi @ptrblck!
One thing I’ve just changed which seems to have solved the missing grad attributes issue is by changing how I defined the shape. So for example, when cycling through optim.param_groups I got the shape via the code below.
module = group['module']
g = module.weight.grad
s = g.shape

However, I’ve changed it to,
g = module['params'][0].grad
s = g.shape

which seems to work which I assume because I’m accessing it via the parameters rather than the module, and the grad attributes are filled when calling loss.backward()?
Thank you!