More robust DCE pass (#4565)

Summary:
X-link: pytorch/pytorch#132764

Pull Request resolved: #4565

- make default DCE pass check schema,

- mark Proxy dump as NotImplemented for better error msg

- Remove Proxy from tensors when dumping models, as Proxy cannot be dumped.

More details in https://docs.google.com/document/d/1G5vmTXjzxoyVGRI2kpA1gQukK_Glyg2NrE0Oh6Nlg9A/edit?usp=sharing.

Reviewed By: angelayi

Differential Revision: D60319175

fbshipit-source-id: d9815ca07e1caaa49bb6a82836765030a67a4270

Author: yushangdi

exir/serde/serialize.py

@@ -342,9 +342,19 @@ def serialize(
             )
         else:
             example_inputs = b""
+
+        # TODO (shangdi): change to use `torch._export.utils.remove_proxy_from_state_dict`
+        # after pytorch repo is updated.
+        # Proxy cannot be dumped, so we remove them.
+        new_state_dict = {}
+        for k, v in exported_program.state_dict.items():
+            if "proxy" in v.__dict__:
+                new_state_dict[k] = v.clone().detach()
+            else:
+                new_state_dict[k] = v
         return export_serialize._SerializedProgram(
             serialized_ep,
-            export_serialize.serialize_torch_artifact(exported_program.state_dict),
+            export_serialize.serialize_torch_artifact(new_state_dict),
             export_serialize.serialize_torch_artifact(constants),
             example_inputs,
         )