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:
brkirch 2022-10-10 00:38:55 -04:00 committed by AUTOMATIC1111
parent 45fbd1c5fe
commit 8acc901ba3
1 changed files with 5 additions and 1 deletions

View File

@ -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':