[cherry-pick] Support FP16 for index_select op (#38751)

* optimize backward (#37055) * update * update * update * modify code style

[cherry-pick] Support FP16 for index_select op (#38751)
* optimize backward (#37055) * update * update * update * modify code style
0d081cbc · Haohongxiang · GitHub · 327f635e · 0d081cbc
隐藏空白更改
内联并排

Showing with 4 addition and 0 deletion

paddle/fluid/operators/index_select_op.cu paddle/fluid/operators/index_select_op.cu +4 -0

未找到文件。
--- a/paddle/fluid/operators/index_select_op.cu
+++ b/paddle/fluid/operators/index_select_op.cu
@@ -201,12 +201,16 @@ REGISTER_OP_CUDA_KERNEL(
    index_select,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, float>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, double>,
+    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext,
+                               paddle::platform::float16>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, int>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, int64_t>);
 REGISTER_OP_CUDA_KERNEL(
    index_select_grad,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext, float>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext, double>,
+    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext,
+                                   paddle::platform::float16>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext, int>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext,
                                   int64_t>);