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.

1 files changed, 5 insertions, 1 deletions

diff --git a/modules/safe.py b/modules/safe.py
index 4d06f2a5..05917463 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -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':