Elocin:

    error = weights[1] + weights[2]



Just a question, why do you have your weights Tensor be of size 64 but only define the error of the first two? (granted they are normalised so there will be gradient for all gradients). I’m not 100% sure but doing weights[1] might be in-place so check your Tensor does indeed have a grad_fn.  Otherwise, your gradient will be zero!
Also, you might want to reduce your learning rate and momentum constant, you’re likely converging directly into a local minima (or perhaps even diverging) and getting stuck. Try lr=1e-4 with/without momentum =0.9.