【报错记录】RuntimeError: don‘t know how to restore data location of torch.FloatStorage (tagged with CPU)

问题描述

在调试一个代码时，在输入参数都正确的情况下，报出 RuntimeError: don't know how to restore data location of torch.FloatStorage (tagged with CPU) 的错误，如下图所示：

解决方案

经过查找资料，在一个外过有人同样的报错记录里找到了答案。

原地址链接：https://github.com/microsoft/DirectML/issues/196

该问题是 torch.load() 的错误，对应到我的代码中，是下面这行代码，只不过它是DML，而我的是CPU，改法都一样：

model.load_state_dict(torch.load(args.pre_train, map_location='CPU'),strict=True)

将其改为：

model.load_state_dict(torch.load(args.pre_train, map_location={
    
    '0':'CPU'}), strict=True)

问题得以解决。

总结

Github上的代码好多都是在Linux上跑的，我们在Windows上跑时，就经常出现类似的问题，将备的参数设置为0貌似是个通用的方法，类似的还有 num_workers 参数的问题。

参考资料

[1] https://github.com/microsoft/DirectML/issues/196