直接在代码中加入下面
import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"
内容源自:https://www.cnblogs.com/darkknightzh/p/6836568.html
pytorch指定使用的单个GPU
原文:https://www.cnblogs.com/xym4869/p/11957993.html