此贴仅为记录debug过程,为防后续再次遇见
问题
问题情境
复现文章模型,使用GPU跑代码,有两张GPU,设置在 cuda: 1 上跑
问题描述
在模型测试加载最优模型时报错:torch.cuda.device_count()显示GPU卡数量只有一张,无法加载使用卡1
报错信息
运行超时
RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.
解决
发现是以下代码问题
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
该语句限制了 PyTorch 进程可以看到的 GPU。例如,如果 x=“0”,则你的代码只能看到一张 GPU(设备 0),即使你机器有多张卡,之后你调用 torch.cuda.device_count() 就只会返回 1
将该行代码注释后可以成功使用GPU 1 运行测试