import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model, device_ids=[2,3])
如果出现 RuntimeError: all tensors must be on devices[0]:
该问题由于PyTorch程序设置的gpu_id编号未从0开始导致的。
解决方案,在运行python命令前设置一下CUDA_VISIBLE_DEVICES:
CUDA_VISIBLE_DEVICES=2,3 python example.py --gpu_id 0 1
注意:在程序中设置gpu_id仍要从0开始,上例中gpu_id 0对应device 2,gpu_id 1对应device 3