KFrank:

Please note that this approach – zeroing out parts of .grad before
calling opt.step() – doesn’t work in general. Some optimizers (e.g.,
when using momentum or weight decay) will change the weights
even if .grad is zero.


Thanks, this is really important note!