/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 vote170 views2 answers
Render template in django
Python Django
0 vote141 views1 answer
Pandas heading
Python
0 vote136 views1 answer
Python - caéculate the average based on the level of a second column in a df
Python
75 votes134 views3 answers
How to get a random element from an array in python
Python
6 votes158 views5 answers
No module named env.__main__; 'env' is a package and cannot be directly executed
Python
1 vote188 views1 answer
Flask sending post request
Python Flask
0 vote141 views1 answer
Python keyboard key codes
Python
4 votes161 views1 answer
How to delete every row in excel using openpyxl
Python
21 votes143 views6 answers
Python remove key from dict
Python
0 vote171 views1 answer
How to save multiple choices in django site:stackoverflow.com
Python Django