From 182b6f8385b17f880652782924e131bbfcc7f9f7 Mon Sep 17 00:00:00 2001 From: Yuanle Liu 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 c6ed50501f766..054964419d2d9 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 bb0fbdf6ca848..6a94e14d7e68b 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; } }