I have been struggling with implementing a custom loss function, and not sure why it is not working properly. If I use one of the provided loss functions, and perform a training step via
loss_fn = nn.MSELoss()
loss = loss_fn(predicted, target)
loss.backward()                                                                                                  
optimizer.step() 

everything works well, and the loss decreases as expected:










Train epoch: 0
Loss: 0.802491


Train epoch: 100
Loss: 0.010990


Train epoch: 200
Loss: 0.009987


Train epoch: 300
Loss: 0.009558


Train epoch: 400
Loss: 0.007305



However, if I change the loss calculation to
loss = torch.mean((predicted - target) ** 2)

it does not work:










Train epoch: 0
Loss: 4.732191


Train epoch: 100
Loss: 0.255895


Train epoch: 200
Loss: 0.256236


Train epoch: 300
Loss: 0.253145


Train epoch: 400
Loss: 0.255263



Any suggestions for where I am going wrong would be greatly appreciated.
Thanks