I have tried to set CUDA_VISIBLE_DEVICES in shell, then I run a simple script to test if the setting has taken effect, unfortunately, it seems does not work.
CUDA_VISIBLE_DEVICES=3; python test.py

the script test.py is
import torch
print(torch.cuda.current_device())

the above script still shows that current device is 0.
I find that torch.cuda.set_device(device_num) works fine in setting the desired GPU to use.