[inference][trt] argmax support int64 dtype (#49167)

fd1730e4 · Zhang Jun · GitHub · baa98d1d · fd1730e4
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

paddle/fluid/inference/tensorrt/op_teller.cc paddle/fluid/inference/tensorrt/op_teller.cc +1 -1

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

    if (op_type == "arg_min") {