Reserve the memory for vector to save cost in gather_ranges_to_dense_op.h (#73540)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/73540

Reviewed By: r-barnes

Differential Revision: D34492500

fbshipit-source-id: f973931b05880c76c82f799d1b40b128078d6702
(cherry picked from commit 64a51de23efaca0d29df88d141aaea40fadddbb0)
This commit is contained in:
Su Shi
2022-03-01 10:12:59 -08:00
committed by PyTorch MergeBot
parent 8ac7393565
commit 1c621a7bfe

View File

@ -146,6 +146,7 @@ class GatherRangesToDenseOp final : public Operator<Context> {
auto& key = Input(KEY);
auto* key_data = key.template data<int64_t>();
vector<std::pair<int64_t, const char*>> buffer;
buffer.reserve(rangeLength);
for (const auto b_i : c10::irange(rangeLength)) {
int64_t one_key_item = key_data[rangeStart + b_i];
auto* one_data_item = rawData + (rangeStart + b_i) * itemsize;