From 499d81e5733b6a80ce28140e9925d2a0c8dc4fcb Mon Sep 17 00:00:00 2001
From: yujianfeng <yujianfeng5@huawei.com>
Date: Wed, 26 Aug 2020 16:02:45 +0800
Subject: [PATCH] Fix remove internal output for unique device target

---
 .../optimizer/ascend/format_type/remove_internal_output.cc  | 2 +-
 .../ascend/format_type/remove_internal_output_test.cc       | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/mindspore/ccsrc/backend/optimizer/ascend/format_type/remove_internal_output.cc b/mindspore/ccsrc/backend/optimizer/ascend/format_type/remove_internal_output.cc
index 530e180ed2e..a7c7ff6a09f 100644
--- a/mindspore/ccsrc/backend/optimizer/ascend/format_type/remove_internal_output.cc
+++ b/mindspore/ccsrc/backend/optimizer/ascend/format_type/remove_internal_output.cc
@@ -58,7 +58,7 @@ const AnfNodePtr RemoveInternalOutput::Process(const FuncGraphPtr &func_graph, c
   if (kernel_graph == nullptr) {
     return nullptr;
   }
-  if (!kernel_graph->IsInternalOutput(node, 0)) {
+  if (!kernel_graph->IsUniqueTargetInternalOutput(node, 0)) {
     return nullptr;
   }
   if (!UsedForOutputOnly(func_graph, node)) {
diff --git a/tests/ut/cpp/pre_activate/ascend/format_type/remove_internal_output_test.cc b/tests/ut/cpp/pre_activate/ascend/format_type/remove_internal_output_test.cc
index d4a69e70a59..a084ff477bd 100644
--- a/tests/ut/cpp/pre_activate/ascend/format_type/remove_internal_output_test.cc
+++ b/tests/ut/cpp/pre_activate/ascend/format_type/remove_internal_output_test.cc
@@ -49,7 +49,7 @@ class TestHWRemoveInternalOutput : public BackendCommon {
     auto make_tuple = GetMakeTuple(kg);
     auto add = make_tuple->cast<CNodePtr>()->input(1);
     MS_EXCEPTION_IF_NULL(add);
-    kg->AddInternalOutput(add, add);
+    kg->AddInternalOutput(add, add, 0, true);
     KernelBuildInfoBuilder builder;
     builder.SetInputsFormat({kOpFormat_DEFAULT, kOpFormat_DEFAULT});
     builder.SetInputsDeviceType({kFloat32->type_id(), kFloat32->type_id()});
@@ -77,8 +77,8 @@ class TestHWRemoveInternalOutput : public BackendCommon {
     MS_EXCEPTION_IF_NULL(tuple_getitem2);
     auto max_pool = tuple_getitem1->cast<CNodePtr>()->input(1);
     MS_EXCEPTION_IF_NULL(max_pool);
-    kg->AddInternalOutput(tuple_getitem1, max_pool);
-    kg->AddInternalOutput(tuple_getitem2, max_pool);
+    kg->AddInternalOutput(tuple_getitem1, max_pool, 0, true);
+    kg->AddInternalOutput(tuple_getitem2, max_pool, 1, true);
     KernelBuildInfoBuilder builder;
     builder.SetInputsReshapeType({{}});
     builder.SetOutputsReshapeType({{}, {}});