search

fastnlp / fastNLP

fastNLP: A Modularized and Extensible NLP Framework. Currently still in incubation.
https://gitee.com/fastnlp/fastNLP
Apache License 2.0
3.07k stars 448 forks source link

在模型加载函数中，增加设定map_location=cpu可选项 #350

Open Yemaoxin opened 3 years ago

Yemaoxin commented 3 years ago

当使用带有GPU的服务跑代码，保存模型后，放在本地没有gpu的电脑进行模型预测时，报错。 RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU. To Reproduce

重现这个bug的步骤 在多gpu服务器训练并保存模型 在本地不带GPU电脑加载时 Expected behavior

增加map_locaiton选项，使得训练的模型可以在本地快速进行新数据的与醋

Desktop (please complete the following information):

yhcc commented 3 years ago

Trainer保存模型的时候已经把模型放置在了cpu才保存的，所以应该不需要map_location就可以了。