/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

0 vote152 views1 answer
How to add a command to button in a loop tkinter python
Python
0 vote155 views1 answer
Install ansible with pip
Python
4 votes125 views3 answers
Module 'tensorflow._api.v2.train' has no attribute 'gradientdescentoptimizer'
Python
1 vote170 views1 answer
Python .nlargest
Python
14 votes163 views2 answers
Python list of random values
Python
0 vote159 views1 answer
Aes in django
Python
6 votes155 views2 answers
Merge list elements python
Python
0 vote154 views2 answers
Print(), .lower() and .upper() are all different _____________ you can use in python.
Python
0 vote178 views1 answer
Nested list flask
Python
0 vote166 views1 answer
Do while python using dates
Python