Solved by Nikronic in post #2 
              
              
                Hi, 
I think the issue is that you are using sum to reduce the MSE loss. 

In nn.MSELoss (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'. 
[image] 
criterion = nn.MSELoss(reduction='s…