stuff related to torch version change

This commit is contained in:
AUTOMATIC
2023-04-29 12:36:50 +03:00
parent 9eb49b04e3
commit ee71eee181
4 changed files with 10 additions and 11 deletions

View File

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