@tom Oops, you’re right. I was following the original implementation on a perceptron too closely.
It seems that the model updates correctly on vanilla Pytorch, but it doesn’t on Pytorch Lightning. If you work with the latter, I made a notebook 1 for you to reproduce the error, but otherwise I will post this in the PL community. Thanks again!