correct conditions of gather in opteller

b3602sss · b3602sss · commit 3c26824b7549 · 2021-07-07T05:36:12.000Z
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -314,8 +314,14 @@ 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<int>("axis")) return false;
     }
 
     if (op_type == "gather_nd") {

Original file line number	Diff line number	Diff line change
`@@ -314,8 +314,14 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,`
`314`	`314`	`}`
`315`	`315`
`316`	`316`	`if (op_type == "gather") {`
	`317`	`+ if (!with_dynamic_shape) return false;`
	`318`	`+`
	`319`	`+ auto inputs = desc.InputArgumentNames();`
	`320`	`+ for (auto& input : inputs) {`
	`321`	`+ if (input == "Axis" && desc.Input("Axis").size() > 0) return false;`
	`322`	`+ }`
`317`	`323`	`// current not support axis from input, use default 0`
`318`		`- if (!with_dynamic_shape \|\| desc.Input("Axis").size() > 0) return false;`
	`324`	`+ if (desc.GetAttrIfExists<int>("axis")) return false;`
`319`	`325`	`}`
`320`	`326`
`321`	`327`	`if (op_type == "gather_nd") {`