模型测试报错：有2张显卡但cuda.device_count()显示GPU卡数量只有一张

此贴仅为记录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 运行测试