fix_arg (#49770)

5d60ff91 · xiaoxiaohehe001 · GitHub · 8cd0d5b3 · 5d60ff91
显示空白变更内容
内联并排

Showing with 6 addition and 2 deletion

paddle/fluid/inference/tensorrt/op_teller.cc paddle/fluid/inference/tensorrt/op_teller.cc +6 -2

未找到文件。
--- a/paddle/fluid/inference/tensorrt/op_teller.cc
+++ b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -686,8 +686,12 @@ struct SimpleOpTypeSetTeller : public Teller {
      int axis = desc.HasAttr("axis")
                     ? PADDLE_GET_CONST(int64_t, desc.GetAttr("axis"))
                     : -1;
-      bool flatten = PADDLE_GET_CONST(bool, desc.GetAttr("flatten"));
+      bool flatten = desc.HasAttr("flatten")
-      int dtype = PADDLE_GET_CONST(int, desc.GetAttr("dtype"));
+                         ? PADDLE_GET_CONST(bool, desc.GetAttr("flatten"))
+                         : false;
+      int dtype = desc.HasAttr("dtype")
+                      ? PADDLE_GET_CONST(int, desc.GetAttr("dtype"))
+                      : 3;
      if (axis == 0 || flatten || (dtype != 2 && dtype != 3)) return false;
    }