From da651f08be0fb0ded8b22313f8ec48de4bccaca2 Mon Sep 17 00:00:00 2001
From: Chen Ding <candy.dc@alibaba-inc.com>
Date: Thu, 20 Jun 2024 14:53:59 +0800
Subject: [PATCH] [Embedding] Fix Incremental EV restore bug.

---
 .../framework/embedding/single_tier_storage.h | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/framework/embedding/single_tier_storage.h b/tensorflow/core/framework/embedding/single_tier_storage.h
index bf3ca4ccc2e..1982ed436bf 100644
--- a/tensorflow/core/framework/embedding/single_tier_storage.h
+++ b/tensorflow/core/framework/embedding/single_tier_storage.h
@@ -85,16 +85,15 @@ class SingleTierStorage : public Storage<K, V> {
 
   void Insert(K key, ValuePtr<V>** value_ptr,
               int64 alloc_len) override {
-    do {
-      *value_ptr = layout_creator_->Create(alloc_, alloc_len);
-      Status s = kv_->Insert(key, *value_ptr);
-      if (s.ok()) {
-        break;
-      } else {
-        (*value_ptr)->Destroy(alloc_);
-        delete *value_ptr;
-      }
-    } while (!(kv_->Lookup(key, value_ptr)).ok());
+    Status s = kv_->Lookup(key, value_ptr);
+    if (s.ok()) {
+      (*value_ptr)->Destroy(alloc_);
+      delete *value_ptr;
+      TF_CHECK_OK(kv_->Remove(key));
+    }
+
+    *value_ptr = layout_creator_->Create(alloc_, alloc_len);
+    TF_CHECK_OK(kv_->Insert(key, *value_ptr));
   }
 
   void Insert(K key, ValuePtr<V>* value_ptr) override {