Merge branch 'dev' into remove-watermark-option

1 files changed, 1 insertions, 4 deletions

diff --git a/modules/safe.py b/modules/safe.py
index 82d44be3..dadf319c 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -1,6 +1,5 @@
 # this code is adapted from the script contributed by anon from /h/

 

-import io

 import pickle

 import collections

 import sys

@@ -12,11 +11,9 @@ import _codecs
 import zipfile

 import re

 

-

 # 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

@@ -27,7 +24,7 @@ class RestrictedUnpickler(pickle.Unpickler):
 

     def persistent_load(self, saved_id):

         assert saved_id[0] == 'storage'

-        return TypedStorage()

+        return TypedStorage(_internal=True)

 

     def find_class(self, module, name):

         if self.extra_handler is not None: