pytorch load函数中cpu与gpu转化

  1. cpu -> cpu或者gpu -> gpu:

checkpoint = torch.load(‘modelparameters.pth’)

model.load_state_dict(checkpoint)
2. cpu -> gpu 1

torch.load(‘modelparameters.pth’, map_location=lambda storage, loc: storage.cuda(1))
3. gpu 1 -> gpu 0

torch.load(‘modelparameters.pth’, map_location={‘cuda:1’:‘cuda:0’})
4. gpu -> cpu

torch.load(‘modelparameters.pth’, map_location=lambda storage, loc: storage)

参见https://blog.csdn.net/qq_41192383/article/details/87871375

你可能感兴趣的:(pytorch)