/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 votes147 views2 answers
Python code fpr fibon
Python
2 votes179 views1 answer
How to convert gb to mb in python
Python
0 vote127 views1 answer
Var colors python
Python
5 votes147 views1 answer
Install pygaame using pip
Python
4 votes139 views2 answers
Code server
Python
6 votes143 views6 answers
Ternary operator in python
Python
0 vote136 views2 answers
Python time a code segment
Python
0 vote165 views1 answer
How to convert datasets into pandas dataframes
Python
0 vote178 views1 answer
Flask where to put db.create_all
Python Flask
0 vote177 views1 answer
Python generator for reading and writing file
Python