/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 vote143 views1 answer
Read tar.gz file python
Python
2 votes118 views6 answers
Variables in python
Python
12 votes102 views3 answers
Setwd python
Python
0 vote109 views1 answer
Pandas groupby histogram
Python
0 vote111 views1 answer
Get current url with parameters url django
Python
0 vote96 views1 answer
Show avg value in sns boxplot
Python
0 vote97 views3 answers
Pandas pivot table
Python
2 votes135 views1 answer
Average out all rows pandas
Python
26 votes113 views2 answers
Dictionary from two list
Python
0 vote102 views1 answer
Minimum Number of Operations to Move All Balls to Each Box in python used in function method
Python