在pycharm中设置环境变量有两种方式，第一种方式是通过右上角Edit Configurations...界面设置：
将cuda_visible_devices设置为5后，将模型model = model.cuda()和输入数据input = input.cuda()加载到gpu上，观察到：
而通过nvidia-smi命令观察到GPU使用情况如图：
可见，设置环境变量CUDA_VISIBLE_DEVICES=5后，系统会对设置的GPU进行重新编号，从’0’开始。cuda()命令会将模型/数据加载到相对“可见”的第一个GPU上。
第二种方式是在代码里设置，也能达到同样的效果：
我看网上有帖子说环境变量os.environ['CUDA_VISIBLE_DEVICES']='5,6'的设置，否则会失效。但我自己试了一下，放到model.cuda()之前都没问题。
当然，如果调换一下顺序，在将模型加载到cuda之后设置环境变量，此时设置会失效，模型参数会放在cuda:0上，后续的input = input.cuda()后，input的device也为cuda:0。
所以只要是在使用cuda之前设置应该都可！（当然为了减少不必要的麻烦，比如说import的其他文件可能会先使用到cuda，还是尽早设置环境变量为好）
参考：