From 182b6f8385b17f880652782924e131bbfcc7f9f7 Mon Sep 17 00:00:00 2001
From: Yuanle Liu <yuanlehome@163.com>
Date: Thu, 27 Apr 2023 19:44:10 +0800
Subject: [PATCH] scale trt converter support int64 (#53388)

---
 paddle/fluid/inference/tensorrt/convert/scale_op.cc | 2 +-
 paddle/fluid/inference/tensorrt/op_teller.cc        | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/scale_op.cc b/paddle/fluid/inference/tensorrt/convert/scale_op.cc
index c6ed50501f7..054964419d2 100644
--- 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
index bb0fbdf6ca8..6a94e14d7e6 100644
--- 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;
         }
       }
-- 
GitLab