From c832625fca29f6463e8710d141ae7025da57e598 Mon Sep 17 00:00:00 2001
From: DefTruth <31974251+DefTruth@users.noreply.github.com>
Date: Thu, 18 Jul 2024 14:41:49 +0800
Subject: [PATCH] Update build.sh

---
 build.sh | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/build.sh b/build.sh
index 7b912af3..d540430b 100755
--- 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