File tree Expand file tree Collapse file tree 1 file changed +4
-5
lines changed Expand file tree Collapse file tree 1 file changed +4
-5
lines changed Original file line number Diff line number Diff line change @@ -128,14 +128,13 @@ TRTEngine::TRTEngine(
128
128
for (size_t pyt_idx = 0 ; pyt_idx < outputs; pyt_idx++) {
129
129
auto binding_name = _out_binding_names[pyt_idx];
130
130
auto trt_idx = cuda_engine->getBindingIndex (binding_name.c_str ());
131
- std::string engine_binded_name = cuda_engine->getIOTensorName (inputs_size + pyt_idx);
132
- TORCHTRT_CHECK (
133
- (binding_name == engine_binded_name),
134
- " Could not find a TensorRT engine binding for output named " << binding_name);
131
+ TORCHTRT_CHECK ((trt_idx != -1 ), " Could not find a TensorRT engine binding for output named " << binding_name);
135
132
TORCHTRT_CHECK (
136
133
!(cuda_engine->getTensorIOMode (binding_name.c_str ()) == nvinfer1::TensorIOMode::kINPUT ),
137
134
" Binding " << binding_name << " specified as output but found as input in TensorRT engine" );
138
- LOG_DEBUG (" Output binding name: " << binding_name << " pyt return idx: " << inputs_size + pyt_idx << " )" );
135
+ LOG_DEBUG (
136
+ " Output binding name: " << binding_name << " has TensorRT binding index: " << trt_idx
137
+ << " , Torch binding index: " << inputs_size + pyt_idx);
139
138
out_binding_map[trt_idx] = pyt_idx;
140
139
out_binding_names[pyt_idx] = binding_name;
141
140
}
You can’t perform that action at this time.
0 commit comments