So while the loss decreases the outputs stay approx. the same? Let me know if that’s the case and I do some more debugging.
One issue that comes to my mind is that we’ve recently had similar problems using nn.MSELoss when the model output and target had a shape mismatch and we’re silently broadcasted. Could you check the shapes of all tensors used in nn.MSELoss?