From d027619b82998d0ff24df122a6a5037af01b7446 Mon Sep 17 00:00:00 2001
From: ling <lingqiaomin@huawei.com>
Date: Tue, 6 Jul 2021 15:09:30 +0800
Subject: [PATCH] support user set graph-output-tensor data

---
 mindspore/lite/src/lite_session.cc    |  5 +++++
 mindspore/lite/src/mindrt_executor.cc | 12 +++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mindspore/lite/src/lite_session.cc b/mindspore/lite/src/lite_session.cc
index 5dfd8f7f865..4911a819b12 100644
--- a/mindspore/lite/src/lite_session.cc
+++ b/mindspore/lite/src/lite_session.cc
@@ -716,6 +716,11 @@ LiteSession::~LiteSession() {
     if (tensor->IsConst() && !tensor->own_data()) {
       tensor->set_data(nullptr);
     }
+
+    /* situation : user set graph-output-tensor data */
+    if (tensor->IsGraphOutput() && tensor->allocator() == nullptr) {
+      tensor->set_data(nullptr);
+    }
     delete tensor;
     tensor = nullptr;
   }
diff --git a/mindspore/lite/src/mindrt_executor.cc b/mindspore/lite/src/mindrt_executor.cc
index 55b2581f88a..b424b7616f8 100644
--- a/mindspore/lite/src/mindrt_executor.cc
+++ b/mindspore/lite/src/mindrt_executor.cc
@@ -121,8 +121,18 @@ void MindrtExecutor::TransferGraphOutput() {
 
 void MindrtExecutor::FreeOutputTensor() {
   for (auto tensor_map : *output_tensor_map_) {
+    auto src_tensor = tensor_map.first;
     auto dst_tensor = tensor_map.second;
-    dst_tensor->FreeData();
+    if (dst_tensor->allocator() != nullptr) {
+      dst_tensor->FreeData();
+    } else {
+      if (dst_tensor->data_type() == src_tensor->data_type()) {
+        /* user set graph-output-tensor from outside */
+        src_tensor->set_data(dst_tensor->data());
+        src_tensor->set_own_data(false);
+        src_tensor->set_allocator(nullptr);
+      }
+    }
   }
   return;
 }