Why setting CUDA_VISIBLE_DEVICES within the code doesn‘t work?

Hi! I’m adding os.environ['CUDA_VISIBLE_DEVICES'] = "2" in my code does not work, the code always select first GPU, however CUDA_VISIBLE_DEVICES=2 python train.py works.
I find that os.environ['CUDA_VISIBLE_DEVICES'] works in this code . I have check my code for multiple times and haven’t got the reason. Who can help me? Thanks a lot.

you need to set that before the first usage of cuda, or maybe even before importing torch.