[ez] fixup the export of D62879819 (#136900)

izaitsevfb · pytorchmergebot · commit b35f70da0555 · 2024-09-28T13:46:17.000Z
a line from D62879819 (#136190) went missing somehow Pull Request resolved: #136900 Approved by: https://github.com/atalman
diff --git a/torch/export/unflatten.py b/torch/export/unflatten.py
@@ -1442,5 +1442,6 @@ def _swap_modules(
         entry.fqn: entry.signature for entry in ep.module_call_graph if entry.signature
     }
     gm = ep.module()
+    gm.graph.eliminate_dead_code()
     assert isinstance(gm, torch.fx.GraphModule)
     return _swap_module_helper(gm, modules_to_swap, module_call_graph)
