如何使用torch.DataParallel

# 先增加可用的cuda环境
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = MyModel()
# 最好不要使用cuda(0)
# 先进行dataparallel的封装，再把模型转移到cuda上面
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3]).to('cuda:0')
或者
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3]).to(device)
# 注意criterion网络也需要在cuda上面
self.criterion = self.criterion.to(device)
# 之后将数据放在同一张gpu上面
data = data.to(device)
target = target.to(device)
output = self.net(batch_sampling_X)
loss = self.criterion(output, batch_src_y)