torch多GPU模型的训练与保存

使用多gpu训练时

model = torch.nn.DataParallel(model, device_ids=[1, 2, 3, 4])

若模型采用多GPU训练，则在模型保存时：

torch.save(model.module.state_dict(), model_out_path)

若单GPU则：

torch.save(mode.state_dict(), model_out_path)