diff --git a/paddle/phi/kernels/gpu/take_along_axis_kernel.cu b/paddle/phi/kernels/gpu/take_along_axis_kernel.cu
index 63113e3e672f37fea9af2bd350a6849a399ff7f1..9665a917d9dc4a5fa0d350f8fb635ecec79b7832 100644
--- a/paddle/phi/kernels/gpu/take_along_axis_kernel.cu
+++ b/paddle/phi/kernels/gpu/take_along_axis_kernel.cu
@@ -53,6 +53,7 @@ PD_REGISTER_KERNEL(take_along_axis,
                    GPU,
                    ALL_LAYOUT,
                    phi::TakeAlongAxisKernel,
+                   float,
                    double,
                    int64_t,
                    int,