fix is_inplace_node check

Differential Revision: D77462717

Pull Request resolved: #12071

Author: JacobSzwejbka

exir/passes/insert_write_back_for_buffers_pass.py

@@ -72,16 +72,15 @@ def _is_inplace_node(node: torch.fx.Node) -> bool:
     """Check if a node is an inplace node."""
     return (
         node.op == "call_function"
-        and isinstance(node.target, torch._ops.OpOverload)
+        and hasattr(node.target, "_schema")
         and is_inplace_variant(
-            node.target._schema.name, node.target._schema.overload_name
+            node.target._schema.name, node.target._schema.overload_name  # pyre-ignore
         )
     )
 
 
 def _inplace_lineage(
     output_arg: torch.fx.Node,
-    gm: torch.fx.GraphModule,
     gs: ExportGraphSignature,
     kind: SchemaKind,
 ) -> bool:
@@ -152,7 +151,6 @@ def insert_write_back_for_buffers_pass(
             # if the arg and target are not the same, we add a copy_ node.
             not _inplace_lineage(
                 output_node.args[0][i],
-                gm,
                 ep.graph_signature,
                 ep.graph_signature.output_specs[i].kind,
             )

exir/tests/test_reinplace_pass.py

@@ -61,6 +61,15 @@ def forward(
 
         self.assertIsNotNone(index_put_node, "Should find an index_put_ node")
 
+        # Find the copy_ node
+        copy_node = None
+        for node in et.exported_program().graph.nodes:
+            if node.op == "call_function" and "copy_" in str(node.target):
+                copy_node = node
+                break
+
+        self.assertIsNone(copy_node, "Shouldn't find an copy_ node")
+
         e = _load_for_executorch_from_buffer(et.buffer)
         self.assertTrue(
             torch.allclose(