/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

-1 vote148 views1 answer
Sort dict based on other list
Python
0 vote152 views1 answer
How to check mix types in pandas column
Python
5 votes111 views1 answer
If string include word in list
Python
0 vote154 views1 answer
Using built-in crawlers is very simple. A minimal example is shown as follows.
Python
1 vote143 views3 answers
For j in range python
Python
0 vote108 views1 answer
Postgres select if data exists
Python
0 vote156 views1 answer
Python pop a element by index
Python
16 votes148 views4 answers
Python remove last instance of a list
Python
0 vote149 views1 answer
Check django channels with redis setup
Python
0 vote142 views1 answer
Google translator api pyhton
Python