pytorch 指定GPU训练
# 1: torch.cuda.set_device(1) # 2: device = torch.device("cuda:1") # 3:(官方推荐)
import os os.environ["CUDA_VISIBLE_DEVICES"] = '1' (同时调用两块GPU的话) os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
- 上一篇 »[转] pytorch指定GPU
- 下一篇 »pytorch GPU训练好的模型使用CPU加载