用nn.DataParallel重新包装一下

device_ids=[0, 1, 2]
model = model.cuda(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)

optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.001)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)
#因为它在DataParallel里面，所以要先变成普通的nn.SGD对象，然后才能调用该类的梯度更新方法。
optimizer.module.step() 

for param_lr in optimizer.module.param_groups: #同样是要加module
    param_lr['lr'] /= 2

def init_criterion():
    criterion = loss.CrossEntropyLoss2d()
    criterion = torch.nn.DataParallel(
            criterion, range(gpu_nums)).cuda()  # range(self.settings.n_gpu)
    return criterion
    
# criterion = init_criterion()
criterion = loss.CrossEntropyLoss2d()

loss = criterion(out, labels_tensor)
loss /= N
optimizer.zero_grad()
# loss.backward()
loss.sum().backward()