Newer versions of PyTorch use TypedStorage instead
Pytorch 1.13 and later will rename _TypedStorage to TypedStorage, so check for TypedStorage and use _TypedStorage if it is not available. Currently this is needed so that nightly builds of PyTorch work correctly.
This commit is contained in:
parent
45fbd1c5fe
commit
8acc901ba3
|
@ -12,6 +12,10 @@ import _codecs
|
||||||
import zipfile
|
import zipfile
|
||||||
|
|
||||||
|
|
||||||
|
# PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage
|
||||||
|
TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage
|
||||||
|
|
||||||
|
|
||||||
def encode(*args):
|
def encode(*args):
|
||||||
out = _codecs.encode(*args)
|
out = _codecs.encode(*args)
|
||||||
return out
|
return out
|
||||||
|
@ -20,7 +24,7 @@ def encode(*args):
|
||||||
class RestrictedUnpickler(pickle.Unpickler):
|
class RestrictedUnpickler(pickle.Unpickler):
|
||||||
def persistent_load(self, saved_id):
|
def persistent_load(self, saved_id):
|
||||||
assert saved_id[0] == 'storage'
|
assert saved_id[0] == 'storage'
|
||||||
return torch.storage._TypedStorage()
|
return TypedStorage()
|
||||||
|
|
||||||
def find_class(self, module, name):
|
def find_class(self, module, name):
|
||||||
if module == 'collections' and name == 'OrderedDict':
|
if module == 'collections' and name == 'OrderedDict':
|
||||||
|
|
Loading…
Reference in New Issue