pytorch在CPU和GPU上加载模型

pytorch允许把在GPU上训练的模型加载到CPU上,也允许把在CPU上训练的模型加载到GPU上。


CPU->CPU,GPU->GPU

torch.load('gen_500000.pkl')


GPU->CPU

torch.load('gen_500000.pkl', map_location=lambda storage, loc: storage)


CPU->GPU1

torch.load('gen_500000.pkl', map_location=lambda storage, loc: storage.cuda(1))



原文地址:https://www.cnblogs.com/mtcnn/p/9411732.html