fix rec trt

5d64077d · andyjpaddle · 81e0e4ba · 5d64077d
显示空白变更内容
内联并排

Showing with 7 addition and 2 deletion

tools/infer/utility.py tools/infer/utility.py +7 -2

未找到文件。
--- 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,
+            if mode == "rec":
-                                              opt_input_shape)
+                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()