The gradients in pytorch can be computed on scalar outputs in pytorch. Your loss is not a scalar as you have used reduction=‘none’. Either use reduction=‘mean’ or reduction=‘sum’ or just calculate mean loss using
loss = loss.mean()

before calling backward.