/Python/Pickle.load from gpu device to cpu

Pickle.load from gpu device to cpu

#if you are using pickle load on a cpu device, from a gpu device
#you will need to override the pickle load functionality like this:
#note: if loading directly in torch, just add map_location='cpu' in load()

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.com

Similar codes

44 votes195 views3 answers
Python string cut last character
Python
1 vote193 views1 answer
Flask send client to another web page
Python
2 votes164 views1 answer
Find the number of nan per column pandas
Python
0 vote155 views1 answer
Django creat app return _bootstrap._gcd_import
Python Django
1 vote165 views1 answer
Count number of repeats in list python
Python
0 vote160 views1 answer
Extend tuple python
Python
0 vote174 views1 answer
Random Remarks Example in python
Python
0 vote194 views1 answer
Check if object is array like python
Python
21 votes204 views2 answers
Getting pi in python
Python
0 vote171 views1 answer
Pyevtk documentation writearraystovtk
Python