From 219d30529aaa33ec00591621992c55beaed7f3c1 Mon Sep 17 00:00:00 2001 From: jianghui58 Date: Sat, 30 Jul 2022 17:09:43 +0800 Subject: [PATCH] fix memory overflow bug --- .../legacy_optimizer/graph/dropout_node_remove_pass.cc | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/mindspore/lite/tools/converter/legacy_optimizer/graph/dropout_node_remove_pass.cc b/mindspore/lite/tools/converter/legacy_optimizer/graph/dropout_node_remove_pass.cc index 87c14bd98bf..1719965299d 100644 --- a/mindspore/lite/tools/converter/legacy_optimizer/graph/dropout_node_remove_pass.cc +++ b/mindspore/lite/tools/converter/legacy_optimizer/graph/dropout_node_remove_pass.cc @@ -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); }); - *matchedTensor = inDataTensorIdx; + if (matchedTensor != gOutTensorIdx.end()) { + *matchedTensor = inDataTensorIdx; + } // find poseNode auto postNodeIdxes = GetOutputNodeIdx(*graphT, nodeIdx, 0); for (auto postNodeIdx : postNodeIdxes) { @@ -71,7 +73,9 @@ 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); }); - *iter = inDataTensorIdx; + if (iter != postNode->inputIndex.end()) { + *iter = inDataTensorIdx; + } } }