From 412877e613a9e01ca1e10db747ae85b7e085c72e Mon Sep 17 00:00:00 2001
From: duanboqiang <firestonelib@gmail.com>
Date: Thu, 26 Aug 2021 11:14:05 +0800
Subject: [PATCH] fix cast op (#35156)

---
 paddle/fluid/operators/cast_op.cc | 1 +
 paddle/fluid/operators/cast_op.cu | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/paddle/fluid/operators/cast_op.cc b/paddle/fluid/operators/cast_op.cc
index 952e9ca329f..5fc97924ef2 100644
--- a/paddle/fluid/operators/cast_op.cc
+++ b/paddle/fluid/operators/cast_op.cc
@@ -119,6 +119,7 @@ REGISTER_OPERATOR(cast, ops::CastOp,
 REGISTER_OP_CPU_KERNEL(
     cast, ops::CastOpKernel<CPU, float>, ops::CastOpKernel<CPU, double>,
     ops::CastOpKernel<CPU, int>, ops::CastOpKernel<CPU, int64_t>,
+    ops::CastOpKernel<CPU, int>, ops::CastOpKernel<CPU, int16_t>,
     ops::CastOpKernel<CPU, bool>, ops::CastOpKernel<CPU, uint8_t>,
     ops::CastOpKernel<CPU, paddle::platform::float16>,
     ops::CastOpKernel<CPU, paddle::platform::bfloat16>,
diff --git a/paddle/fluid/operators/cast_op.cu b/paddle/fluid/operators/cast_op.cu
index 0beb2291060..8b2998e52a1 100644
--- a/paddle/fluid/operators/cast_op.cu
+++ b/paddle/fluid/operators/cast_op.cu
@@ -102,6 +102,7 @@ REGISTER_OP_CUDA_KERNEL(
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, double>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, int>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, int64_t>,
+    ops::CastOpKernel<paddle::platform::CUDADeviceContext, int16_t>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, bool>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, uint8_t>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext,
@@ -116,6 +117,7 @@ REGISTER_OP_CUDA_KERNEL(
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, double>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, int>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, int64_t>,
+    ops::CastOpKernel<paddle::platform::CUDADeviceContext, int16_t>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, bool>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext, uint8_t>,
     ops::CastOpKernel<paddle::platform::CUDADeviceContext,
-- 
GitLab