In my case (on a cluster with gpus), adding CUDA_VISIBLE_DEVICES=GPU_ID before the python command solved the problem. Instead of adding CUDA_VISIBLE_DEVICES in the command line, you can probably just add a line os.environ[“CUDA_VISIBLE_DEVICES”] =“GPU_ID” before import torch.
2 Likes