Don’t use in-place operations like “+=”. They don’t play nicely with autograd, you can read more here 2.
So instead of:
loss += (1 - torch.pow(param, 2)).sum()
use:
loss = loss + (1 - torch.pow(param, 2)).sum()
Try this change and see if it affects what you’re seeing.