/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

130 votes98 views10 answers
Conda create environment
Python
7 votes114 views2 answers
Find max length of list of list python
Python
0 vote121 views1 answer
How to read numbers from a text file in python
Python
0 vote134 views1 answer
Python tensorflow is not defined
Python
5 votes160 views1 answer
How to sort a list in python using lambda
Python
1 vote135 views1 answer
Error: Command '['/home/robert/python/python_p/env/bin/python3.8', '-Im', 'ensurepip', '--upgrade', '--default-pip']' returned non-zero exit status 1.
Python
15 votes127 views6 answers
Django models
Python
10 votes115 views3 answers
Pandas create new column conditional on other columns
Python
0 vote136 views3 answers
Python convert dat file to csv
Python
4 votes142 views2 answers
Python replace by dictionary
Python