From 246d7e4fd8455f5ed5332b434b82d85b67f15358 Mon Sep 17 00:00:00 2001
From: y00500818 <yangyongqiang7@hisilicon.com>
Date: Thu, 17 Jun 2021 10:36:12 +0800
Subject: [PATCH] bugfix for restore context

---
 ge/generator/ge_generator.cc | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/ge/generator/ge_generator.cc b/ge/generator/ge_generator.cc
index 575afb35..58047c89 100644
--- a/ge/generator/ge_generator.cc
+++ b/ge/generator/ge_generator.cc
@@ -674,6 +674,12 @@ Status GeGenerator::GenerateModel(const Graph &graph, const string &file_name_pr
     GELOGD("Current ctx is null.");
     ctx = nullptr;
   }
+  std::function<void()> callback = [&]() {
+    if (ctx != nullptr) {
+      (void)rtCtxSetCurrent(ctx);
+    }
+  };
+  GE_MAKE_GUARD(restore, callback);
 
   GeRootModelPtr ge_root_model = nullptr;
   GE_CHECK_NOTNULL_EXEC(impl_, return PARAM_INVALID);
@@ -712,11 +718,6 @@ Status GeGenerator::GenerateModel(const Graph &graph, const string &file_name_pr
     }
     return ret;
   }
-
-  if (ctx != nullptr) {
-    (void)rtCtxSetCurrent(ctx);
-  }
-
   return SUCCESS;
 }