Remove .1 from FXOutputName so it matches the fx graph node name

Xilinx · Jul 28, 2023 · 57ed408 · 57ed408
1 parent 364dd51
commit 57ed408
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp b/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp
@@ -90,7 +90,11 @@ void NodeImporter::importNode(Node *node, MlirBlock appendToBlock,
 
   auto addFxOutputNameAttr = [&](MlirOperation& operation) {
     if (output && output->hasDebugName()) {
-      auto strAttr = mlirStringAttrGet(context, toMlirStringRef(output->debugName()));
+      std::string name = output->debugName();
+      size_t len = name.size();
+      if (len > 2 && name[len-2] == '.' && name[len-1] == '1')
+        name = name.substr(0, len-2);
+      auto strAttr = mlirStringAttrGet(context, toMlirStringRef(name));
       mlirOperationSetAttributeByName(operation, toMlirStringRef("FXOutputName"), strAttr);
     }
   };