Hi,
Why do you want add  .mean() on loss function, if you call the method in torch.nn or torch.nn.functional it has already call .mean(). For example, if you use torch.nn.CrossEntropyLoss() as your loss function, and input two tensors, one is your model’s output, the other is correspond target, and the loss is also a tensor with only one element like this.
tensor(1.6975, device='cuda:1', grad_fn=<NllLoss2DBackward>)
... ...

And you can do backpropagation in your network normally.