shape op support int8 and uint8 tensor (#30201)

4 years ago · 609c022222
parent 01a287bf0a
commit 609c022222
2 changed files with 3 additions and 1 deletions
--- a/paddle/fluid/operators/shape_op.cc
+++ b/paddle/fluid/operators/shape_op.cc
@ -69,5 +69,6 @@ REGISTER_OPERATOR(
    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
 REGISTER_OP_CPU_KERNEL(shape, ops::ShapeKernel<bool>, ops::ShapeKernel<int>,
+                       ops::ShapeKernel<int8_t>, ops::ShapeKernel<uint8_t>,
                       ops::ShapeKernel<int64_t>, ops::ShapeKernel<float>,
                       ops::ShapeKernel<double>);
--- a/paddle/fluid/operators/shape_op.cu
+++ b/paddle/fluid/operators/shape_op.cu
@ -16,7 +16,8 @@ limitations under the License. */

 REGISTER_OP_CUDA_KERNEL(
    shape, paddle::operators::ShapeKernel<bool>,
-    paddle::operators::ShapeKernel<int>,
+    paddle::operators::ShapeKernel<int>, paddle::operators::ShapeKernel<int8_t>,
+    paddle::operators::ShapeKernel<uint8_t>,
    paddle::operators::ShapeKernel<int64_t>,
    paddle::operators::ShapeKernel<float>,
    paddle::operators::ShapeKernel<double>,