If I don’t choose GPU, everything works fine. The thing is I had to share GPUs with others so I had to choose GPUs to run.

Sure, but I was suggesting to do that via “device” not “os.environ”. It’s a bit weird because os.environ[“CUDA_VISIBLE_DEVICES”] = “1”) is supposed to set the default device to GPU 1 but via “cuda:0” you are using GPU 0. I would just try to only use “device” and not mix it via
os.environ[“CUDA_VISIBLE_DEVICES”] = “1”)

because I suspect that if you use os.environ[“CUDA_VISIBLE_DEVICES”] = “1”), then “cuda:0”  will effectively be  the first card that is visible, i.e., GPU 1, which can easily become confusing and then lead to issues.
I suggest setting
DEVICE = torch.device("cuda:1")

in your script and then use DEVICE everywhere, like
b = torch.tensor(arr, dtype=torch.int64, device=DEVICE)

and don’t use this
os.environ[“CUDA_VISIBLE_DEVICES”] = “1”

because it’s not necessary then.