Fix DeepStream 5 and 6

2024-11-13 10:10:12 -03:00
parent b451b036b2
commit 1b96a01ec3
1 changed files with 5 additions and 0 deletions
--- a/nvdsinfer_custom_impl_Yolo/yolo.cpp
+++ b/nvdsinfer_custom_impl_Yolo/yolo.cpp
@@ -193,7 +193,12 @@ Yolo::createEngine(nvinfer1::IBuilder* builder)
  config->setProfilingVerbosity(nvinfer1::ProfilingVerbosity::kDETAILED);
 #endif

+#if NV_TENSORRT_MAJOR >= 8 && NV_TENSORRT_MINOR > 0
  nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(*builder->getLogger());
+#else
+  nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
+#endif
+
  assert(runtime);

  nvinfer1::IHostMemory* serializedEngine = builder->buildSerializedNetwork(*network, *config);