Pytorch FAQ

1、pytorch指定使用哪个GPU

使用CUDA_VISIBLE_DEVICES，类似tensorflow指定GPU的方法。 直接终端中设定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

python代码中设定：

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

使用函数set_device，该函数见 pytorch-master\torch\cuda\__init__.py。

import torchtorch.cuda.set_device(id)

不过官方建议使用CUDA_VISIBLE_DEVICES，不建议使用set_device函数。