[fix] Fixed buffer_size setting when hkv loads and saves files

tensorflow · Apr 30, 2024 · 697341f · 697341f
1 parent eabccef
commit 697341f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/...flow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv.h b/...flow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_hkv.h
@@ -565,7 +565,9 @@ class TableWrapper {
       throw std::runtime_error(error_msg);
     }
 
-    size_t n_saved = table_->save(wfile.get(), buffer_size, stream);
+    size_t n_saved = table_->save(
+        wfile.get(),
+        buffer_size * (sizeof(K) + sizeof(V) * dim + sizeof(uint64_t)), stream);
     if (has_scores) {
       LOG(INFO) << "[op] Save " << n_saved << " pairs from keyfile: " << keyfile
                 << ", and valuefile: " << valuefile << ", and scorefile"
@@ -616,7 +618,9 @@ class TableWrapper {
       throw std::runtime_error(error_msg);
     }
 
-    size_t n_loaded = table_->load(rfile.get(), buffer_size, stream);
+    size_t n_loaded = table_->load(
+        rfile.get(),
+        buffer_size * (sizeof(K) + sizeof(V) * dim + sizeof(uint64_t)), stream);
     if (has_scores) {
       LOG(INFO) << "[op] Load " << n_loaded
                 << " pairs from keyfile: " << keyfile