fix rec trt

5d64077d · andyjpaddle · 81e0e4ba · 5d64077d
Commit 5d64077d authored Dec 06, 2021 by andyjpaddle
Hide whitespace changes
Inline Side-by-side

Showing with 7 additions and 2 deletions

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

No files found.
--- 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()