Error: grad can be implicitly created only for scalar outputs

[ptrblck]

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