>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
torch.load(f, map_location=None, pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, **pickle_load_args)
参数 | 描述 |
---|---|
f | |
map_location | |
pickle_module |
参考:
https://pytorch.org/docs/stable/generated/torch.load.html#torch.load