diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index f945f5166190..0be88a1e1ba6 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/deepspeed/runtime/zero/parameter_offload.py @@ -52,7 +52,7 @@ def __init__(self, parent_module, *args, **kwargs): def __reduce__(self): r0, _, *r2 = super().__reduce__() - return (r0, (self._parent_module, )) + r2 + return (r0, (self._parent_module, )) + tuple(r2) def __getitem__(self, key): param = super().__getitem__(key)