diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc index cfbe3957ad444f116182fb396b64964da06213a8..f98b0c9ede76e2ec542a0c1d74ea13d0201e57f9 100644 --- a/paddle/fluid/inference/tensorrt/op_teller.cc +++ b/paddle/fluid/inference/tensorrt/op_teller.cc @@ -314,8 +314,13 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8, } if (op_type == "gather") { + if (!with_dynamic_shape) return false; + auto inputs = desc.InputArgumentNames(); + for (auto& input : inputs) { + if (input == "Axis" && desc.Input("Axis").size() > 0) return false; + } // current not support axis from input, use default 0 - if (!with_dynamic_shape || desc.Input("Axis").size() > 0) return false; + if (desc.GetAttrIfExists("axis")) return false; } if (op_type == "gather_nd") {