This commit is contained in:
Won-Kyu Park 2025-12-20 01:16:16 +01:00 committed by GitHub
commit 5324465f44
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -37,6 +37,9 @@ class RestrictedUnpickler(pickle.Unpickler):
if res is not None:
return res
class Empty:
pass
if module == 'collections' and name == 'OrderedDict':
return getattr(collections, name)
if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']:
@ -51,12 +54,8 @@ class RestrictedUnpickler(pickle.Unpickler):
return getattr(numpy, name)
if module == '_codecs' and name == 'encode':
return encode
if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint':
import pytorch_lightning.callbacks
return pytorch_lightning.callbacks.model_checkpoint
if module == "pytorch_lightning.callbacks.model_checkpoint" and name == 'ModelCheckpoint':
import pytorch_lightning.callbacks.model_checkpoint
return pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint
if module.startswith("pytorch_lightning"):
return Empty
if module == "__builtin__" and name == 'set':
return set