Error(s) in loading state_dict for Module

经过分析是用torch0.4.1版本训练，用0.4.0版本预测的时候报的异常or

查看torch版本

import torch
print(torch.__version__)  #注意是双下划线

经过多次测试，我发现，这个问题是由以下原因引起的：

当你训练的时候用多卡训练，但是测试的时候用一张卡测试，或者反过来，就会出现类似于上图这种error

将

model.load_state_dict(checkpoint['state_dict'])

改为

model.load_state_dict({k.replace('module.',''):v for k,v in torch.load(checkpoint_path)['state_dict'].items()})

即可