From 5d64077deafad455e2d3c62763726bbfdb52bdb8 Mon Sep 17 00:00:00 2001 From: andyjpaddle Date: Mon, 6 Dec 2021 10:18:21 +0000 Subject: [PATCH] fix rec trt --- tools/infer/utility.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index af50c5e6..d7e058c2 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -271,8 +271,13 @@ def create_predictor(args, mode, logger): min_input_shape = {"x": [1, 3, 10, 10]} max_input_shape = {"x": [1, 3, 512, 512]} opt_input_shape = {"x": [1, 3, 256, 256]} - config.set_trt_dynamic_shape_info(min_input_shape, max_input_shape, - opt_input_shape) + if mode == "rec": + if args.rec_algorithm == "CRNN": + config.set_trt_dynamic_shape_info( + min_input_shape, max_input_shape, opt_input_shape) + else: + config.set_trt_dynamic_shape_info( + min_input_shape, max_input_shape, opt_input_shape) else: config.disable_gpu() -- GitLab