fix(tensorrt): update tensorrt code of traffic_light_classifier (auto…

…warefoundation#2325) Signed-off-by: M. Fatih Cırıt <[email protected]> Signed-off-by: kminoda <[email protected]>
kminoda · Jan 6, 2023 · a544503 · a544503
1 parent c2e2e24
commit a544503
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/perception/traffic_light_classifier/utils/trt_common.cpp b/perception/traffic_light_classifier/utils/trt_common.cpp
@@ -75,8 +75,16 @@ void TrtCommon::setup()
   }
 
   context_ = UniquePtr<nvinfer1::IExecutionContext>(engine_->createExecutionContext());
-  input_dims_ = engine_->getBindingDimensions(getInputBindingIndex());
-  output_dims_ = engine_->getBindingDimensions(getOutputBindingIndex());
+
+#if (NV_TENSORRT_MAJOR * 10000) + (NV_TENSORRT_MINOR * 100) + NV_TENSOR_PATCH >= 80500
+  input_dims_ = engine_->getTensorShape(input_name_.c_str());
+  output_dims_ = engine_->getTensorShape(output_name_.c_str());
+#else
+  // Deprecated since 8.5
+  input_dims_ = engine_->getBindingDimensions(engine_->getBindingIndex(input_name_.c_str()));
+  output_dims_ = engine_->getBindingDimensions(engine_->getBindingIndex(output_name_.c_str()));
+#endif
+
   is_initialized_ = true;
 }
 
@@ -155,8 +163,4 @@ int TrtCommon::getNumOutput()
     output_dims_.d, output_dims_.d + output_dims_.nbDims, 1, std::multiplies<int>());
 }
 
-int TrtCommon::getInputBindingIndex() { return engine_->getBindingIndex(input_name_.c_str()); }
-
-int TrtCommon::getOutputBindingIndex() { return engine_->getBindingIndex(output_name_.c_str()); }
-
 }  // namespace Tn
diff --git a/perception/traffic_light_classifier/utils/trt_common.hpp b/perception/traffic_light_classifier/utils/trt_common.hpp
@@ -118,8 +118,6 @@ class TrtCommon
   bool isInitialized();
   int getNumInput();
   int getNumOutput();
-  int getInputBindingIndex();
-  int getOutputBindingIndex();
 
   UniquePtr<nvinfer1::IExecutionContext> context_;