There might be some other tensors in the computation graph which are in different device.
There are a few ways to avoid this. please see Which device is model / tensor stored on?
One way to check what device your tensors are in is : tensor.device, you can check to see if all tensors are in the same device.
You only need to set device for your inputs and models, the other tensors which are created during the computation will be automatically in the same device. e.g. loss.