ptrblck:

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.


That is exactly the case. While the loss function is decreasing the outputs remain the same.



 ptrblck:

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 ?


OK! I checked, and the shapes are not the same!
output = torch.Size([2, 2])
target = torch.Size([2])

Why stack is changing the shape of the output tensor??? What can I use instead of stack? I really appreciate all your help!!