RuntimeError: Input tensors need to be on the same GPU, but found the following tensor and device combinations:

[kristosh]

When I am trying to run the notebook on linear probing which runs smoothly after installing your package transformers-heads, I am getting into the following issue in the the like regarding training:

RuntimeError: Input tensors need to be on the same GPU, but found the following tensor and device combinations:
 [(torch.Size([1, 884736]), device(type='cuda', index=1)), (torch.Size([27648]), device(type='cuda', index=0)), (torch.Size([2304, 768]), device(type='cuda', index=1))]

How can I cope with this issue in my case and why it is happening for me?