Hi,
I think the issue is that you are using sum to reduce the MSE loss.



 Pranavan_Theivendira:

return torch.sum((inputs - target) ** 2)


In nn.MSELoss 1 (please see the linked documentation), you can specify an argument reduction which by default is mean. So, you may need to explicitly state you want reduction='sum'.

criterion = nn.MSELoss(reduction='sum')

Bests