From 5787fdb830398c2579b468630b424b16e03cf8ba Mon Sep 17 00:00:00 2001
From: kswang <wangkaisheng2@huawei.com>
Date: Thu, 5 Nov 2020 10:18:43 +0800
Subject: [PATCH] set optimizer flag for assign

---
 mindspore/ccsrc/backend/session/kernel_graph.cc | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/mindspore/ccsrc/backend/session/kernel_graph.cc b/mindspore/ccsrc/backend/session/kernel_graph.cc
index 2a57707a7d..07c8176771 100644
--- a/mindspore/ccsrc/backend/session/kernel_graph.cc
+++ b/mindspore/ccsrc/backend/session/kernel_graph.cc
@@ -1226,10 +1226,20 @@ void KernelGraph::SetOptimizerFlag() {
   has_optimizer_ = false;
   for (const auto &cnode : execution_order_) {
     MS_EXCEPTION_IF_NULL(cnode);
-    if (kOptOperatorSet.find(AnfAlgo::GetCNodeName(cnode)) != kOptOperatorSet.end()) {
+    auto node_name = AnfAlgo::GetCNodeName(cnode);
+    if (kOptOperatorSet.find(node_name) != kOptOperatorSet.end()) {
       has_optimizer_ = true;
       return;
     }
+    if (node_name.find("Assign") != string::npos) {
+      for (auto &input : cnode->inputs()) {
+        MS_EXCEPTION_IF_NULL(input);
+        if (input->isa<Parameter>() && AnfAlgo::IsParameterWeight(input->cast<ParameterPtr>())) {
+          has_optimizer_ = true;
+          return;
+        }
+      }
+    }
   }
 }