RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'target' in call to _thnn_nll_loss_forward

vision
7

Thanks for the update.
The error is caused, since you are explicitly casting the targets to torch.LongTensor, which is a CPU type.
Instead of using .type(torch.LongTensor), call label1.squeeze().long() to keep the tensor on the GPU.

1 Like