cfr

tensorflow · Mar 31, 2024 · 3f4a9d2 · 3f4a9d2
1 parent 1b94e7a
commit 3f4a9d2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/...rflow_recommenders_addons/dynamic_embedding/python/kernel_tests/hkv_hashtable_ops_test.py b/...rflow_recommenders_addons/dynamic_embedding/python/kernel_tests/hkv_hashtable_ops_test.py
@@ -713,7 +713,8 @@ def test_hkv_hashtable_save_and_load_all_with_local_file_system(self):
           initializer=init_ops.random_normal_initializer(0.0, 0.01),
           devices=devices,
           dim=dim,
-          kv_creator=de.HkvHashTableCreator(),
+          kv_creator=de.HkvHashTableCreator(config=de.HkvHashTableConfig(
+              max_capacity=99999)),
       )
       var2 = de.get_variable(
           'lfslav2_' + str(idx),
@@ -722,7 +723,8 @@ def test_hkv_hashtable_save_and_load_all_with_local_file_system(self):
           initializer=init_ops.random_normal_initializer(0.0, 0.01),
           devices=devices,
           dim=dim,
-          kv_creator=de.HkvHashTableCreator(),
+          kv_creator=de.HkvHashTableCreator(config=de.HkvHashTableConfig(
+              max_capacity=99999)),
       )
       init_keys = tf.range(0, 20000, dtype=tf.int64)
       init_values = var1.lookup(init_keys)