Fix to custom op lowering transform

llvm · Jun 13, 2022 · a48a6a4 · a48a6a4
1 parent 357a335
commit a48a6a4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/lib/Conversion/TorchToLinalg/CustomOpExample.cpp b/lib/Conversion/TorchToLinalg/CustomOpExample.cpp
@@ -37,8 +37,7 @@ class ConvertCustomOpExample
       return failure();
     // Since the example op does nothing, we simply replace the uses of the
     // return value with its argument, then remove the op.
-    op->replaceAllUsesWith(op->getOperands());
-    op->erase();
+    rewriter.replaceOp(op, op->getOperands());
 
     return success();
   }

diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/registry.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/registry.py
@@ -139,7 +139,7 @@ def uppercase_first_letter(s):
         cpp_class_name = "".join(
             uppercase_first_letter(s) for s in op_class_name_atoms) + "Op"
         # Disallow leading underscores in C++ to avoid illegal names.
-        cpp_class_name = cpp_class_name.lstrip('_')
+        cpp_class_name = cpp_class_name.lstrip("_")
         return op_name, cpp_class_name
 
     def get_shape_function_signature(self):