scale trt converter support int64 (#53388)

PaddlePaddle · Apr 27, 2023 · 182b6f8 · 182b6f8
1 parent 1c97aa6
commit 182b6f8
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/paddle/fluid/inference/tensorrt/convert/scale_op.cc b/paddle/fluid/inference/tensorrt/convert/scale_op.cc
@@ -19,7 +19,7 @@ namespace inference {
 namespace tensorrt {
 
 /*
- * ConcatOp
+ * Scale Op
  */
 class ScaleOpConverter : public OpConverter {
  public:

diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -1184,10 +1184,12 @@ struct SimpleOpTypeSetTeller : public Teller {
           return false;
         }
       } else {
-        // At present, only support float32 or float16 or int32 into trt.
+        // At present, only support float32 or float16 or int32 or int64 into
+        // trt.
         if (!(dtype == framework::proto::VarType::FP32 ||
               dtype == framework::proto::VarType::FP16 ||
-              dtype == framework::proto::VarType::INT32)) {
+              dtype == framework::proto::VarType::INT32 ||
+              dtype == framework::proto::VarType::INT64)) {
           return false;
         }
       }