pytorch多GPU训练时,报错Broadcast function not implemented for CPU tensors
通过DataParallel设置pytorch的多GPU训练时,出现错误 “Broadcast function not implemented for CPU tensors”
进行如下设置即可解决.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
# optimizer = nn.DataParallel(optimizer)
model.to(device) #加入此句,可以解决相应问题