From a2a12fe459b181c5018f9bc528cd298f1c1dfb22 Mon Sep 17 00:00:00 2001
From: LDOUBLEV <liuvv0203@outlook.com>
Date: Thu, 25 Nov 2021 09:42:37 +0800
Subject: [PATCH] fix TRT8 core bug

---
 tools/infer/utility.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tools/infer/utility.py b/tools/infer/utility.py
index 41a3c0f1..15e5abdf 100755
--- a/tools/infer/utility.py
+++ b/tools/infer/utility.py
@@ -173,6 +173,7 @@ def create_predictor(args, mode, logger):
         config.enable_use_gpu(args.gpu_mem, 0)
         if args.use_tensorrt:
             config.enable_tensorrt_engine(
+                workspace_size=1 << 30,
                 precision_mode=precision,
                 max_batch_size=args.max_batch_size,
                 min_subgraph_size=args.min_subgraph_size)
@@ -294,7 +295,10 @@ def create_predictor(args, mode, logger):
 
 def get_infer_gpuid():
     cmd = "nvidia-smi"
-    res = os.popen(cmd).readlines()
+    try:
+        res = os.popen(cmd).readlines()
+    except:
+        res = None
     if len(res) == 0:
         return None
     cmd = "env | grep CUDA_VISIBLE_DEVICES"
-- 
GitLab