/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

2 votes221 views1 answer
Python split string at index
Python
0 vote280 views2 answers
Split paragraphs in python
Python
-2 votes263 views1 answer
Numpy difference between two arrays
Python
0 vote238 views1 answer
Python keyborad back space
Python
6 votes276 views1 answer
Jupyter pdataframe max rows
Python
0 vote250 views1 answer
Concatenar columnas en una del mismo dataset
Python
0 vote252 views1 answer
Aggregation with f() in django rest api
Python
2 votes260 views3 answers
How to get an input into a list python
Python
0 vote264 views2 answers
How to get value from txtbox in flask
Python
0 vote246 views2 answers
Pandas agg rename columns
Python