/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

19 votes294 views5 answers
Python floor
Python
0 vote296 views1 answer
IQR to remove outlier
Python
0 vote315 views1 answer
Are there learning activities for django-debug-toolbar
Python Django
0 vote306 views1 answer
Exec to return a value python
Python
2 votes334 views3 answers
Django login required decorator
Python Django
18 votes297 views2 answers
Absolute value python
Python
1 vote320 views1 answer
Django wait for database
Python Django
1 vote327 views1 answer
Open a filename starting with in python
Python
0 vote265 views1 answer
How to make py file open in current directory
Python
1 vote393 views1 answer
Next iteration python
Python