diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc index 052d17878a5a9dfb3c2b2cae4644b3d4dda2942f..72338bcef1197a0051052c653f726e15d576cbd0 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") {