/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

0 vote195 views1 answer
Pandas add prefix to column names
Python
18 votes184 views4 answers
Create a venv
Python
0 vote207 views1 answer
Boto3 cross region
Python
0 vote248 views1 answer
How to extract field values in list from queryset in django
Python Django
0 vote218 views1 answer
Matplotlib giving error "OverflowError: In draw_path: Exceeded cell block limit"
Python
2 votes233 views2 answers
Or operator django
Python Django
0 vote189 views2 answers
1007 solution python
Python
2 votes245 views2 answers
Remove first item from list python
Python
1 vote189 views2 answers
Django print query
Python
3 votes240 views2 answers
Check iterable python
Python