Saving and loading torch models on 2 machines with different number of GPU devices

2

Hi,

Personally, I always save my models on cpu to be able to load them easily anywhere, and put them on gpu later if needed.

However, there is an option in torch.load to “remap storages to be loaded on a different device”.
See this post explaining it: Loading weights for CPU model while trained on GPU
And the doc: http://pytorch.org/docs/master/torch.html?highlight=load#torch.load

Quoting from doc:

torch.load('tensors.pt')
# Load all tensors onto the CPU
torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Map tensors from GPU 1 to GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
4 Likes