!19460 support user set graph-output-tensor data

Merge pull request !19460 from ling/fix
This commit is contained in:
i-robot 2021-07-06 10:59:58 +00:00 committed by Gitee
commit ff9fa4f34d
2 changed files with 16 additions and 1 deletions

View File

@ -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;
}

View File

@ -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;
}