Add __getstate__ and __setstate__ methods to frontend torch.nn.Module

ivy-llc · Dec 2, 2023 · 0984dae · 0984dae
1 parent 4577f47
commit 0984dae
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/ivy/functional/frontends/torch/nn/modules/module.py b/ivy/functional/frontends/torch/nn/modules/module.py
@@ -312,3 +312,11 @@ def __dir__(self):
         keys = [key for key in keys if not key[0].isdigit()]
 
         return sorted(keys)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop("_compiled_call_impl", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)