fix TRT8 core bug

This commit is contained in:
LDOUBLEV
2021-11-25 09:42:37 +08:00
parent fa1a12563a
commit a2a12fe459
+5 -1
View File
@@ -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"