Update build.sh

DefTruth · Jul 18, 2024 · c832625 · c832625
1 parent acb73d1
commit c832625
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/build.sh b/build.sh
@@ -10,6 +10,22 @@ else
 fi
 
 cd "${BUILD_DIR}" && pwd 
-cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel -DCMAKE_INSTALL_PREFIX=./install
+if [ $1 == "tensorrt" ]; then
+  cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
+           -DCMAKE_INSTALL_PREFIX=./install \
+           -DENABLE_TENSORRT=ON \
+           -DCUDA_DIR=/usr/local/cuda \
+           -DTensorRT_DIR=/usr/local/tensorrt \
+           -DENABLE_TEST=ON
+
+else
+  cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \
+           -DCMAKE_INSTALL_PREFIX=./install \
+           -DENABLE_TEST=ON
+fi
+
 make -j8
 make install
+
+# bash ./build.sh
+# bash ./build.sh tensorrt