1、先设置当前环境可见的GPU
import os
os.environ["CUDA_VISIBLE_DEVICES"]="2,3"#如希望使用2,3号GPU
2、创建设备
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")#使用上面设定的全部GPU
#device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#若用单块GPU,指定cuda:0时对应上面设定的2号GPU
3、将模型并行化
model=My_Model()
if torch.cuda.device_count()>1:#多CPU时
model=nn.DataParallel(model)
model.to(device)