踩过的pytorch坑
2021-05-12 本文已影响0人
顾北向南
1. 多卡训练模型
- 如果使用torch.nn.DataParallel(model)多卡并行训练模型的话需要注意:
model = torch.nn.DataParallel(model).module # 必须这样设置模型,不加module的话,模型训练会出现“module must have its parameters and buffers on device cuda:0 but found one o”