Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!39254 [MS][LITE][r1.5] fix heap buffer overflow bug 

Merge pull request !39254 from jianghui58/test_r1.5
This commit is contained in:
i-robot 2022-08-02 02:09:21 +00:00 committed by Gitee
parent dfa9b65f61 219d30529a
commit b69bd64948
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mindspore/lite/tools/converter/legacy_optimizer/graph/dropout_node_remove_pass.cc
View File

@ -1,5 +1,5 @@
/**
* Copyright 2020-2021 Huawei Technologies Co., Ltd
* Copyright 2020-2022 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@ -62,7 +62,9 @@ STATUS IsolateDropoutNode(schema::MetaGraphT *graphT, size_t nodeIdx) {
auto matchedTensor =
std::find_if(gOutTensorIdx.begin(), gOutTensorIdx.end(),
[&outDataTensorIdx](const unsigned int &idx) { return (idx == outDataTensorIdx); });
if (matchedTensor != gOutTensorIdx.end()) {
*matchedTensor = inDataTensorIdx;
}
// find poseNode
auto postNodeIdxes = GetOutputNodeIdx(*graphT, nodeIdx, 0);
for (auto postNodeIdx : postNodeIdxes) {
@ -71,9 +73,11 @@ STATUS IsolateDropoutNode(schema::MetaGraphT *graphT, size_t nodeIdx) {
MS_ASSERT(postNode != nullptr);
auto iter = std::find_if(postNode->inputIndex.begin(), postNode->inputIndex.end(),
[&outDataTensorIdx](const unsigned int &idx) { return (idx == outDataTensorIdx); });
if (iter != postNode->inputIndex.end()) {
*iter = inDataTensorIdx;
}
}
}
// now all node's outputTensors are useless
// remove all node's outputTensors