From 9a480204bd180a6b1636702ab335fd437b0f05b7 Mon Sep 17 00:00:00 2001 From: liangtao07 Date: Wed, 19 Jan 2022 11:32:58 +0800 Subject: [PATCH] fix the bug caused by vector.insert in torch_tensorrt::core::ir::get_value_first_calc_dtype_opt Signed-off-by: liangtao07 --- core/ir/ir.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/core/ir/ir.cpp b/core/ir/ir.cpp index 7da08d5e19..1c1813ea5f 100644 --- a/core/ir/ir.cpp +++ b/core/ir/ir.cpp @@ -58,8 +58,8 @@ c10::optional get_value_first_calc_dtype_opt(torch::jit::Block* auto consumers = in->uses(); auto search_list = std::vector(consumers.begin(), consumers.end()); - for (auto& u : search_list) { - auto n = u.user; + for (auto iter = search_list.begin(); iter != search_list.end(); ++iter) { + auto n = iter->user; LOG_GRAPH("Node we are looking at: " << util::node_info(n)); auto ins = n->inputs(); auto outs = n->outputs(); @@ -94,7 +94,9 @@ c10::optional get_value_first_calc_dtype_opt(torch::jit::Block* for (auto o : outs) { if (o->type() == c10::TensorType::get()) { auto o_uses = o->uses(); + int offset = iter - search_list.begin(); search_list.insert(search_list.end(), o_uses.begin(), o_uses.end()); + iter = search_list.begin() + offset; } } continue; @@ -123,7 +125,9 @@ c10::optional get_value_first_calc_dtype_opt(torch::jit::Block* for (auto o : outs) { if (o->type() == c10::TensorType::get()) { auto o_uses = o->uses(); + int offset = iter - search_list.begin(); search_list.insert(search_list.end(), o_uses.begin(), o_uses.end()); + iter = search_list.begin() + offset; } } } @@ -150,4 +154,4 @@ TypeMap get_block_first_calc_dtypes_opt(torch::jit::Block* b) { } // namespace ir } // namespace core -} // namespace torch_tensorrt \ No newline at end of file +} // namespace torch_tensorrt