Add PADDLE_THROW in take_along_axis kernel when the datatype of index is wrong. (#53556)

paddle/phi/kernels/gpu/take_along_axis_grad_kernel.cu

@@ -30,11 +30,6 @@ void TakeAlongAxisGradKernel(const Context& dev_ctx,
                              const DenseTensor& out_grad,
                              int axis,
                              DenseTensor* x_grad) {
-  PADDLE_ENFORCE_EQ(
-      dev_ctx.GetPlace().GetType() == phi::AllocationType::GPU,
-      true,
-      errors::PreconditionNotMet("This kernel only runs on GPU."));
-
   // We need to know the shape of input matrix to determine the shape of grad
   // matrix of input.
   x_grad->Resize(x.dims());
@@ -55,6 +50,11 @@ void TakeAlongAxisGradKernel(const Context& dev_ctx,
   } else if (index_type == DataType::INT64) {
     phi::funcs::gpu_scatter_add_kernel<T, int64_t>(
         *x_grad, axis, index, out_grad, dev_ctx);
+  } else {
+    PADDLE_THROW(
+        phi::errors::InvalidArgument("The data type of input index is expected "
+                                     "to be int32 or int64, but recieved %s.",
+                                     phi::DataTypeToString(index_type)));
   }
 }
 

paddle/phi/kernels/gpu/take_along_axis_kernel.cu

@@ -28,11 +28,6 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
                          const DenseTensor& index,
                          int axis,
                          DenseTensor* out) {
-  PADDLE_ENFORCE_EQ(
-      dev_ctx.GetPlace().GetType() == phi::AllocationType::GPU,
-      true,
-      errors::PreconditionNotMet("This kernel only runs on GPU device."));
-
   out->Resize(index.dims());
   dev_ctx.template Alloc<T>(out);
 
@@ -41,6 +36,11 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
     phi::funcs::gpu_gather_kernel<T, int32_t>(x, axis, index, *out, dev_ctx);
   } else if (index_type == DataType::INT64) {
     phi::funcs::gpu_gather_kernel<T, int64_t>(x, axis, index, *out, dev_ctx);
+  } else {
+    PADDLE_THROW(
+        phi::errors::InvalidArgument("The data type of input index is expected "
+                                     "to be int32 or int64, but recieved %s.",
+                                     phi::DataTypeToString(index_type)));
   }
 }