/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

5 votes173 views6 answers
Dict comprehension python
Python
1 vote165 views2 answers
How to 404 custom page not found in django
Python Django
0 vote129 views3 answers
How to correct spelling in pandas datafeame
Python
0 vote146 views2 answers
Arithmetic in python
Python
0 vote139 views1 answer
How to swirtch the placement of the levels in pandas
Python
3 votes165 views2 answers
Print first word of a string python and return it
Python
0 vote146 views1 answer
Pandas turn counts into probability
Python
0 vote141 views1 answer
Python is not set from command line or npm configuration node-gyp
Python
3 votes171 views2 answers
Python get current mouse position
Python
3 votes139 views1 answer
How to set datetime format in python
Python