Skip to content

Commit 1e0d282

Browse files
authored
fix: correct tensor deduplication logic (#844)
1 parent fd693ac commit 1e0d282

File tree

1 file changed

+17
-4
lines changed

1 file changed

+17
-4
lines changed

model.cpp

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include <stdarg.h>
2+
#include <algorithm>
23
#include <atomic>
34
#include <chrono>
45
#include <fstream>
@@ -2006,13 +2007,25 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
20062007
w.join();
20072008
}
20082009

2009-
std::unordered_map<std::string, IndexedStorage> latest_map;
2010+
std::vector<IndexedStorage> deduplicated;
2011+
deduplicated.reserve(all_results.size());
2012+
std::unordered_map<std::string, size_t> name_to_pos;
20102013
for (auto& entry : all_results) {
2011-
latest_map[entry.ts.name] = entry;
2014+
auto it = name_to_pos.find(entry.ts.name);
2015+
if (it == name_to_pos.end()) {
2016+
name_to_pos.emplace(entry.ts.name, deduplicated.size());
2017+
deduplicated.push_back(entry);
2018+
} else if (deduplicated[it->second].index < entry.index) {
2019+
deduplicated[it->second] = entry;
2020+
}
20122021
}
20132022

2014-
processed_tensor_storages.reserve(latest_map.size());
2015-
for (auto& [name, entry] : latest_map) {
2023+
std::sort(deduplicated.begin(), deduplicated.end(), [](const IndexedStorage& a, const IndexedStorage& b) {
2024+
return a.index < b.index;
2025+
});
2026+
2027+
processed_tensor_storages.reserve(deduplicated.size());
2028+
for (auto& entry : deduplicated) {
20162029
processed_tensor_storages.push_back(entry.ts);
20172030
}
20182031
}

0 commit comments

Comments
 (0)