load pickle file obtained from GPU to CPU

Question

I meet a problem when I load a pickle file to CPU. I search it on the internet, and they say I need to add map_location parameter. However, after I add this parameter, the problem still exists.

the code is as follows:

torch.__version__
torch.load('featurs.pkl',map_location='cpu')

>>>

'1.0.1.post2'
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.

I know it is because of the different devices, but I use the instruction in the error message, so I do not know how to solve it in the next step.

Thanks in advance!

Answer 1

Try this:

    torch.load('featurs.pkl',map_location=torch.device('cpu'))

Answer 2

The error message suggests using map_location=torch.device('cpu') but even that doesn't work. One workaround is to use the pickle library and implement a custom unpickler.

import pickle
import io

class CPU_Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        else:
        return super().find_class(module, name)

#contents = pickle.load(f) becomes...
contents = CPU_Unpickler(f).load()

Source: Github