From 6af2a247fdb46899e5783a5eb4f3fa69bff15f9f Mon Sep 17 00:00:00 2001
From: zhaozhixuan <zhaozhixuan2@hisilicon.com>
Date: Tue, 23 Mar 2021 13:51:41 +0800
Subject: [PATCH] Add ut.

---
 ge/generator/ge_generator.cc                   | 2 +-
 tests/ut/ge/generator/ge_generator_unittest.cc | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/ge/generator/ge_generator.cc b/ge/generator/ge_generator.cc
index 2ff0c327..65ae5501 100644
--- a/ge/generator/ge_generator.cc
+++ b/ge/generator/ge_generator.cc
@@ -714,7 +714,7 @@ void GeGenerator::RemoveConst(const vector<GeTensor> &inputs, vector<GeTensor> &
   for (auto input : inputs) {
     GeTensorDesc input_desc = input.GetTensorDesc();
     bool is_const = false;
-    (void)AttrUtils::GetBool(tensor, CONST_ATTR_NAME_INPUT, is_const);
+    (void)AttrUtils::GetBool(input_desc, CONST_ATTR_NAME_INPUT, is_const);
     if (!is_const) {
       outputs.emplace_back(input);
     }
diff --git a/tests/ut/ge/generator/ge_generator_unittest.cc b/tests/ut/ge/generator/ge_generator_unittest.cc
index 6d0db429..7b087e94 100644
--- a/tests/ut/ge/generator/ge_generator_unittest.cc
+++ b/tests/ut/ge/generator/ge_generator_unittest.cc
@@ -128,4 +128,13 @@ TEST_F(UtestGeGenerator, test_set_model_name) {
   ge_root_model->root_graph_ = std::move(graph);
   EXPECT_EQ(generator.SetModelNameForDump(ge_root_model), SUCCESS);
 }
+
+TEST_F(UtestGeGenerator, test_remove_const) {
+  GeGenerator generator;
+  GeTensorDesc tensor_desc;
+  GeTensor tensor(tensor_desc);
+  const vector<GeTensor> inputs = {tensor};
+  vector<GeTensor> outputs;
+  generator.RemoveConst(inputs, outputs);
+}
 }  // namespace ge