diff --git a/ge/graph/build/memory/graph_mem_assigner.cc b/ge/graph/build/memory/graph_mem_assigner.cc
index a2e30b52..7b37d4cf 100755
--- a/ge/graph/build/memory/graph_mem_assigner.cc
+++ b/ge/graph/build/memory/graph_mem_assigner.cc
@@ -2126,7 +2126,7 @@ void GraphMemoryAssigner::UpdateMemBlockVisitInfo(const NodePtr &update_node,
                                                   size_t matched_mem_offset,
                                                   map<size_t, pair<NodePtr, vector<int64_t>>> &mem_block_visit_info) {
   if (!is_same_input) {
-    mem_block_visit_info[matched_mem_offset].first = node;
+    mem_block_visit_info[matched_mem_offset].first = update_node;
     mem_block_visit_info[matched_mem_offset].second.clear();
   }
   mem_block_visit_info[matched_mem_offset].second.push_back(in_data_anchor->GetIdx());