From 2d7ed5ffcb77c11c480cead5e141fe74bb10fdea Mon Sep 17 00:00:00 2001
From: acmore <hezj07+github@gmail.com>
Date: Tue, 7 Feb 2023 12:01:41 +0800
Subject: [PATCH] Fix bug that rehash might not be triggered

---
 .../core/kernels/lookup_impl/lookup_table_op_gpu.h              | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h b/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h
index 3b9e2f3cc..60123caba 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h
+++ b/tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h
@@ -173,7 +173,7 @@ class TableWrapper final : public TableWrapperBase<K, V> {
                 << (float)total_size / (float)new_capacity << "].";
     }
     CUDA_CHECK(cudaStreamSynchronize(stream));
-    new_hint_size = table_->get_size(stream);
+    new_hint_size = table_->get_size(stream) + new_keys_num;
     CUDA_CHECK(cudaStreamSynchronize(stream));
     return new_hint_size;
   }