/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 votes180 views2 answers
Django admin actions
Python Django
0 vote221 views1 answer
Flask print request headers
Python Flask
2 votes191 views1 answer
Vowel and consonant list python
Python
13 votes203 views1 answer
Root bg tkinter
Python
12 votes210 views5 answers
Python random uuid
Python
0 vote203 views1 answer
Pyserial read
Python
2 votes211 views1 answer
Open file in python directory
Python
18 votes212 views6 answers
Get last element of array python
Python
7 votes171 views1 answer
Data frame do nympy
Python
-1 vote221 views1 answer
Django migrate not creating tables
Python Django