关于pytorch使用torch.nn.parallel.DistributedDataParallel() 显卡分配不均匀

加载预训练模型默认在第0块gpu上,将其修改为cpu的预训练加载方式:

model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)

你可能感兴趣的:(踩坑系列,pytorch,深度学习,python)