/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

4 votes175 views2 answers
Python get index of substring in liast
Python
0 vote156 views1 answer
Start Openvino Python Application at Boot Time using System Service on ubunut
Python
0 vote466 views1 answer
Reading a list in python
Python
81 votes189 views6 answers
Python convert string to date
Python
3 votes178 views1 answer
Remove commas from string python
Python
1 vote187 views1 answer
How to make a bot say hello <username> when a user says hello in discord with python
Python
0 vote166 views1 answer
Create a typo with python
Python
1 vote170 views1 answer
Use chrome console in selenium
Python
0 vote225 views1 answer
Spacegoat meaning
Python
1 vote192 views1 answer
Set time complexity python
Python