I had this same issue where setting CUDA_VISIBLE_DEVICES=2 python train.py works but setting  os.environ['CUDA_VISIBLE_DEVICES'] = "2" didn’t. The cause of the issue for me was importing the torch packages before setting os.environ['CUDA_VISIBLE_DEVICES'], moving it to the top of the file before importing torch solved it. Hope this helps.