Could you try to run your code with CUDA_VISIBLE_DEVICES=2 python your_script.py or alternatively try to set torch.cuda.set_device(2) right after import torch.
It might be that still some CUDA context is initialized on the default GPU (GPU0) if you try to push your model to another device.