Error: grad can be implicitly created only for scalar outputs

You could try to use reduction='mean' which would lower the loss value or just remove the reduction argument, as mean is the default.

1 Like