/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 vote177 views1 answer
Selenium click on item in a list
Python
1 vote192 views3 answers
Count function in python
Python
0 vote322 views1 answer
Request.args.get check if defined
Python
0 vote196 views1 answer
Pyplot savefig
Python
0 vote205 views1 answer
Round to 3 significant figures python
Python
2 votes258 views3 answers
Play mp3 file python
Python
0 vote185 views1 answer
How to output index of list python without braquets
Python
-1 vote188 views1 answer
Python graph
Python
0 vote201 views1 answer
Additon of multiple duration timestamps python
Python
87 votes304 views2 answers
How to write a while statement in python
Python