From c489f04bb8e0279e249cf1e3599ed0cd4478bb4b Mon Sep 17 00:00:00 2001 From: buxue Date: Fri, 21 Aug 2020 17:15:36 +0800 Subject: [PATCH] fix bug for PynativeExecutor clear. --- mindspore/ccsrc/pipeline/pynative/pynative_execute.cc | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index e9ac967ba1..c2c8d0937e 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -1252,9 +1252,14 @@ void PynativeExecutor::GradNetInner(const GradOperationPtr &grad, const py::obje void PynativeExecutor::Clear(const std::string &flag) { if (!flag.empty()) { MS_LOG(DEBUG) << "Clear res"; - (void)graph_map_.erase(flag); - (void)cell_graph_map_.erase(flag); - (void)cell_resource_map_.erase(flag); + auto key_value = std::find_if(graph_map_.begin(), graph_map_.end(), + [&flag](const auto &item) { return item.first.find(flag) != std::string::npos; }); + if (key_value != graph_map_.end()) { + std::string key = key_value->first; + (void)graph_map_.erase(key); + (void)cell_graph_map_.erase(key); + (void)cell_resource_map_.erase(key); + } Clean(); // Maybe exit in the pynative runing op, so need reset pynative flag. auto ms_context = MsContext::GetInstance();