pytorch使用指定显卡

在调用显卡运行代码时，如果多人使用一张显卡，会遇到out of memory的情况，需要换显卡，但是torch默认使用第0张显卡。
如果在命令窗口运行代码：

python xxx.py

可以在前边加上export CUDA_VISIBLE_DEVICES=0 #这里是要使用的GPU编号，正常的话是从0开始。
如果使用的是pycharm远程连接服务器，可以直接在pycharm中添加环境变量：路径为Run->Edit Configurations->Environment->Envirotment Variable中写入CUDA_VISIBLE_DEVICES=1，如下图所示。
亲测可用。

还有一些别的方法，比如在代码中写入torch.cuda.set_device(0)，尝试了一下，还是失败了，不知道什么原因，依然显示显存不足。