[Embedding] Improve code format.

Signed-off-by: JunqiHu <[email protected]>
DeepRec-AI · Oct 19, 2023 · 796fb20 · 796fb20
1 parent eb68cfd
commit 796fb20
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/tensorflow/core/framework/embedding/cpu_hash_map_kv.h b/tensorflow/core/framework/embedding/cpu_hash_map_kv.h
@@ -148,7 +148,8 @@ class LocklessHashMap : public KVInterface<K, V> {
     for (int64 j = 0; j < bucket_count; j++) {
       if (hash_map_dump[j].first != LocklessHashMap<K, V>::EMPTY_KEY_ 
           && hash_map_dump[j].first != LocklessHashMap<K, V>::DELETED_KEY_
-          && hash_map_dump[j].first % 1000 % partition_nums != partition_id) {
+          && hash_map_dump[j].first % kSavedPartitionNum 
+              % partition_nums != partition_id) {
         key_list->emplace_back(hash_map_dump[j].first);
         value_ptr_list->emplace_back(hash_map_dump[j].second);
       }

diff --git a/tensorflow/core/framework/embedding/dense_hash_map_kv.h b/tensorflow/core/framework/embedding/dense_hash_map_kv.h
@@ -131,7 +131,7 @@ class DenseHashMap : public KVInterface<K, V> {
     }
     for (int i = 0; i< partition_num_; i++) {
       for (const auto it : hash_map_dump[i].hash_map) {
-        if (it.first % 1000 % partition_nums != partition_id) {
+        if (it.first % kSavedPartitionNum % partition_nums != partition_id) {
           key_list->push_back(it.first);
           value_ptr_list->push_back(it.second);
         }

diff --git a/tensorflow/core/framework/embedding/leveldb_kv.h b/tensorflow/core/framework/embedding/leveldb_kv.h
@@ -203,7 +203,7 @@ class LevelDBKV : public KVInterface<K, V> {
     for (it->SeekToFirst(); it->Valid(); it->Next()) {
       K key;
       memcpy((char*)&key, it->key().ToString().data(), sizeof(K));
-      if (key % 1000 % partition_nums == partition_id) continue;
+      if (key % kSavedPartitionNum % partition_nums == partition_id) continue;
       key_list->emplace_back(key);
       FeatureDescriptor<V> hbm_feat_desc(
           1, 1, ev_allocator()/*useless*/,