add missing tensorrt lib search path

open-mmlab · Apr 14, 2022 · 4332b53 · 4332b53
1 parent 823e08b
commit 4332b53
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cmake/modules/FindTENSORRT.cmake b/cmake/modules/FindTENSORRT.cmake
@@ -22,7 +22,7 @@ set(__TENSORRT_LIB_COMPONENTS nvinfer;nvinfer_plugin)
 foreach(__component ${__TENSORRT_LIB_COMPONENTS})
     find_library(
         __component_path ${__component}
-        HINTS ${TENSORRT_DIR}
+        HINTS ${TENSORRT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
         PATH_SUFFIXES lib lib64 lib/x64)
     if (NOT __component_path)
         message(FATAL_ERROR "Cannot find TensorRT lib ${__component} in "