Spaces:
Runtime error
Runtime error
Elvan Selvano
commited on
Commit
β’
02757a2
1
Parent(s):
565d6fb
Upload cpu_unpickler.py
Browse files- cpu_unpickler.py +14 -0
cpu_unpickler.py
ADDED
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import io
|
2 |
+
import pickle
|
3 |
+
|
4 |
+
import torch
|
5 |
+
|
6 |
+
class cpu_unpickler(pickle.Unpickler):
|
7 |
+
"""
|
8 |
+
Overrides the default behavior of the `Unpickler` class to load
|
9 |
+
a `torch.storage` object from abyte string
|
10 |
+
"""
|
11 |
+
def find_class(self, module, name):
|
12 |
+
if module == 'torch.storage' and name == '_load_from_bytes':
|
13 |
+
return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
|
14 |
+
return super().find_class(module, name)
|