pytorch，令代码在指定gpu上运行

所用服务器有四个卡0、1、2、3，想指定代码在2卡上运行。不知道为啥，用

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'

和

export CUDA_VISIBLE_DEVICES= 2

最后用nvidia-smi看代码都在0卡上跑。。
最终解决：

import torch
torch.cuda.set_device(1)

• 我刚开始改成只用torch，会提示我设备号不存在，加上os之后可以执行，后来又试了几次删掉os，也可以执行？？？费解