使用一组特定的 GPU 设备

要使用特定的 GPU 设备集，可以使用 CUDA_VISIBLE_DEVICES 环境变量：

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="0" # Will use only the first GPU device

os.environ["CUDA_VISIBLE_DEVICES"]="0,3" # Will use only the first and the fourth GPU devices

（引自这个答案 ; 这里有关于 CUDA 环境变量的更多信息。）