Hello, @SimonW. Thank you for your reply.
So I changed the codes to be as follows.
import torch
 
if __name__ == '__main__':
    try:
        t = torch.randn(1, device='cuda')
        print('Success!')
    except Exception as e:
        print(e)

It still said OOM on GPU 2.
$ CUDA_VISIBLE_DEVICES=2 python3 check.py 
CUDA error: out of memory