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

View File

@ -12,6 +12,10 @@ import _codecs
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):
out = _codecs.encode(*args)
return out
@ -20,7 +24,7 @@ def encode(*args):
class RestrictedUnpickler(pickle.Unpickler):
def persistent_load(self, saved_id):
assert saved_id[0] == 'storage'
return torch.storage._TypedStorage()
return TypedStorage()
def find_class(self, module, name):
if module == 'collections' and name == 'OrderedDict':