Commit cca16373 authored by brkirch's avatar brkirch

Add attributes used by MPS

parent 16b4509f
...@@ -37,16 +37,16 @@ class RestrictedUnpickler(pickle.Unpickler): ...@@ -37,16 +37,16 @@ class RestrictedUnpickler(pickle.Unpickler):
if module == 'collections' and name == 'OrderedDict': if module == 'collections' and name == 'OrderedDict':
return getattr(collections, name) return getattr(collections, name)
if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter']: if module == 'torch._utils' and name in ['_rebuild_tensor_v2', '_rebuild_parameter', '_rebuild_device_tensor_from_numpy']:
return getattr(torch._utils, name) return getattr(torch._utils, name)
if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage']: if module == 'torch' and name in ['FloatStorage', 'HalfStorage', 'IntStorage', 'LongStorage', 'DoubleStorage', 'ByteStorage', 'float32']:
return getattr(torch, name) return getattr(torch, name)
if module == 'torch.nn.modules.container' and name in ['ParameterDict']: if module == 'torch.nn.modules.container' and name in ['ParameterDict']:
return getattr(torch.nn.modules.container, name) return getattr(torch.nn.modules.container, name)
if module == 'numpy.core.multiarray' and name == 'scalar': if module == 'numpy.core.multiarray' and name in ['scalar', '_reconstruct']:
return numpy.core.multiarray.scalar return getattr(numpy.core.multiarray, name)
if module == 'numpy' and name == 'dtype': if module == 'numpy' and name in ['dtype', 'ndarray']:
return numpy.dtype return getattr(numpy, name)
if module == '_codecs' and name == 'encode': if module == '_codecs' and name == 'encode':
return encode return encode
if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint': if module == "pytorch_lightning.callbacks" and name == 'model_checkpoint':
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment