Skip to content

Commit 23b3fef

Browse files
authored
Make 'scale_op' supporting int and int64 (#5986)
* Make 'scale_op' supporting int and int64 * refine .cu file
1 parent 6375c8c commit 23b3fef

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

paddle/operators/scale_op.cc

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,4 +77,6 @@ REGISTER_OPERATOR(scale, ops::ScaleOp, ops::ScaleOpMaker<float>,
7777
ops::ScaleGradMaker);
7878
REGISTER_OP_CPU_KERNEL(scale,
7979
ops::ScaleKernel<paddle::platform::CPUPlace, float>,
80-
ops::ScaleKernel<paddle::platform::CPUPlace, double>);
80+
ops::ScaleKernel<paddle::platform::CPUPlace, double>,
81+
ops::ScaleKernel<paddle::platform::CPUPlace, int>,
82+
ops::ScaleKernel<paddle::platform::CPUPlace, int64_t>);

paddle/operators/scale_op.cu

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,6 @@
1616

1717
REGISTER_OP_GPU_KERNEL(
1818
scale, paddle::operators::ScaleKernel<paddle::platform::GPUPlace, float>,
19-
paddle::operators::ScaleKernel<paddle::platform::GPUPlace, double>);
19+
paddle::operators::ScaleKernel<paddle::platform::GPUPlace, double>,
20+
paddle::operators::ScaleKernel<paddle::platform::GPUPlace, int>,
21+
paddle::operators::ScaleKernel<paddle::platform::GPUPlace, int64_t>);

0 commit comments

Comments
 (0)