pytorch上配置使用双显卡或多显卡

pytorch上配置使用双显卡

首先那你要有两块或多块显卡，其次torch.cuda.is_available()要返回为true

// 可见设置，环境变量使得指定设备对CUDA应用可见
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
//选择第一块显卡开始载入数据
device = torch.device("cuda: 0")
model.to(device)
//显卡大于1块时，device_ids选择模型载入数据对应的显卡
if torch.cuda.device_count() > 1:
	model= nn.DataParallel(model,device_ids=[0,1])