From 4ea342786528b95c31135afd411c6bd81e89298b Mon Sep 17 00:00:00 2001
From: zlsh80826 <rewang@nvidia.com>
Date: Thu, 18 Mar 2021 14:11:41 +0800
Subject: [PATCH] [Paddle-TRT] support batch axis concatenation when using
 dynamic shape (#31627)

* support batch axis concatenation when using dynamic shape

* opteller can't return true early, or some test will not be executed
---
 paddle/fluid/inference/tensorrt/op_teller.cc | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
index 052d17878a..72338bcef1 100644
--- a/paddle/fluid/inference/tensorrt/op_teller.cc
+++ b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -159,7 +159,11 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,
         return false;
       } else {
         int axis = BOOST_GET_CONST(int, desc.GetAttr("axis"));
-        if (axis <= 0) return false;
+        if (with_dynamic_shape) {
+          if (axis < 0) return false;
+        } else {
+          if (axis <= 0) return false;
+        }
       }
     }
     if (op_type == "transpose2" || op_type == "transpose") {
-- 
GitLab