Speed up _extract_graph_with_inputs_outputs

ghstack-source-id: 6787796ed5033d8bf22730703fbd547575b1a366 Pull Request resolved: #125937
pytorch · May 10, 2024 · 00f811d · 00f811d
1 parent 946b96f
commit 00f811d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch/_functorch/partitioners.py b/torch/_functorch/partitioners.py
@@ -96,7 +96,7 @@ def _extract_graph_with_inputs_outputs(joint_graph, inputs, outputs):
         env[node] = new_node
 
     for node in joint_graph.nodes:
-        if node in inputs:
+        if node in env:
             continue
         elif node.op == "placeholder":
             env[node] = InvalidNode