训练时有多个GPU,或者指定某一个GPU,单和测试的GPU数量不同会报错
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.
解决办法:
指定GPU
torch.load(model_path, map_location = 'cuda:0')
设置当前环境存在的GPU