correct conditions of gather in opteller (#33999)

b3602sss · web-flow · commit 11f5a400b91f · 2021-07-08T20:48:18.000+08:00
* correct conditions of gather in opteller

* test=develop

* test=allcase
diff --git 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<int>("axis")) return false;
     }
 
     if (op_type == "gather_nd") {

Original file line number	Diff line number	Diff line change
`@@ -314,8 +314,13 @@ 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`	`+ auto inputs = desc.InputArgumentNames();`
	`319`	`+ for (auto& input : inputs) {`
	`320`	`+ if (input == "Axis" && desc.Input("Axis").size() > 0) return false;`
	`321`	`+ }`
`317`	`322`	`// current not support axis from input, use default 0`
`318`		`- if (!with_dynamic_shape \|\| desc.Input("Axis").size() > 0) return false;`
	`323`	`+ if (desc.GetAttrIfExists<int>("axis")) return false;`
`319`	`324`	`}`
`320`	`325`
`321`	`326`	`if (op_type == "gather_nd") {`