/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 vote159 views1 answer
Quando è stata inventata la lavastoviglie
Python
0 vote145 views1 answer
Python string: indexing and slicing string
Python
0 vote135 views1 answer
Roganrola
Python
1 vote165 views1 answer
How to loop over day name in python
Python
0 vote305 views1 answer
OduleNotFoundError: No module named 'absl'
Python
1 vote168 views1 answer
Get every item but the last item of python list
Python
0 vote151 views1 answer
How to plot side by side bar horizontal bar graph in python
Python
4 votes208 views6 answers
Python eval
Python
0 vote139 views1 answer
Torchvision.transforms flatten
Python
0 vote172 views1 answer
Check this id exist in database django
Python Django