Skip to content

Commit 6e4aa0b

Browse files
authored
fix: Repair output binding indexing scheme in TRT (#2054)
1 parent f957cc8 commit 6e4aa0b

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

core/runtime/execute_engine.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -155,8 +155,10 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
155155
std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->output_profile_path);
156156
}
157157

158-
for (size_t o = inputs.size(); o < (compiled_engine->num_io.first + compiled_engine->num_io.second); o++) {
159-
uint64_t pyt_idx = compiled_engine->out_binding_map[o];
158+
for (auto output_indices : compiled_engine->out_binding_map) {
159+
// out_binding_map stores TRT_IDX: PYT_IDX
160+
auto pyt_idx = output_indices.second;
161+
160162
std::string name = compiled_engine->out_binding_names[pyt_idx];
161163
auto out_shape = compiled_engine->exec_ctx->getTensorShape(name.c_str());
162164
LOG_DEBUG("Output Name: " << name << " Shape: " << out_shape);

0 commit comments

Comments
 (0)