If you use the same weight you will have to do the for loop I’m afraid.
Note that in your code, you want to do weights.grad.clone() on the last line. All changes to .grad are inplace and so after you do weight.zero_grad(), your buffer will contain only zeros if you don’t clone.