报错:RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1

报错:RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1
原因:在使用Pytorch加载模型时报错。加载的模型是用两个GPU训练的,而加载模型的电脑只有一个GPU,所以会出错。
解决:model = torch.load(model_path)
改为:model = torch.load(model_path, map_location='cuda:0')
如果是4块到2块:就把map_location改为:map_location={'cuda:1': 'cuda:0'}

posted @ 2021-09-02 15:48  小筱痕  阅读(473)  评论(0编辑  收藏  举报