🚀 Feature

The ability to pickle.load a Python object containing a torch cuda tensor on a CPU only machine.

Motivation

Currently, trying to do this gives RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location='cpu' to map your storages to the CPU. Even though you are loading with pickle.load, not torch.load.

When the "loading" code is pytorch agnostic (exists in a repo that does not use pytorch), you can't just change a pickle.load(f) into a torch.load(f, map_location='cpu'). This may be the case when the saved data takes a particular structure and loading/unloading is handled by some code that does not depend on pytorch.

Pitch

A context manager could take care of this:

with torch.loading_context(map_location='cpu'):
    obj = pickle.load(f)  # In my case this call is buried deeper in torch-agnostic code