Can not assign gpu id, help?

[brisker]

I tried to assign my multi gpus with :
CUDA_VISIBLE_DIVICES=2 python train.py ,
but why pytorch is always using the first gpu?

[miguelvr (Miguel Varela Ramos)]

torch.cuda.set_device(gpu_id) # gpu_id is an integer corresponding to your GPU ID