[Fix] fixup the save and load test with new evict strategy.

NVIDIA-Merlin · Oct 11, 2023 · 2cb38b6 · 2cb38b6
1 parent 1e2cd40
commit 2cb38b6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/include/merlin_hashtable.cuh b/include/merlin_hashtable.cuh
@@ -1978,7 +1978,7 @@ class HashTable {
       }
 
       EvictStrategy::set_global_epoch(static_cast<S>(IGNORED_GLOBAL_EPOCH));
-      insert_or_assign(count, d_keys, d_values, d_scores, stream, true);
+      insert_or_assign(count, d_keys, d_values, d_scores, stream, true, true);
       total_count += count;
 
       // Read next batch.

diff --git a/tests/save_and_load_test.cc.cu b/tests/save_and_load_test.cc.cu
@@ -79,8 +79,8 @@ void test_save_to_file() {
                    Table::evict_strategy == EvictStrategy::kEpochLru)
                       ? nullptr
                       : d_scores;
-  table_0->insert_or_assign(keynum, d_keys, d_vectors, temp_score, stream,
-                            global_epoch);
+  EvictStrategy::set_global_epoch(global_epoch);
+  table_0->insert_or_assign(keynum, d_keys, d_vectors, temp_score, stream);
   printf("Fill table_0.\n");
   nv::merlin::LocalKVFile<K, V, S> file;
   std::string keys_path = prefix + ".keys";