[Bug fix]: fix memset on default stream (PaddlePaddle#71696)

Author: zhangyuqin1998

paddle/phi/kernels/gpu/put_along_axis_grad_kernel.cu

@@ -20,6 +20,7 @@
 #include "paddle/phi/core/tensor_utils.h"
 #include "paddle/phi/core/utils/data_type.h"
 #include "paddle/phi/kernels/funcs/gather_scatter_functor.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
 
@@ -91,12 +92,13 @@ void PutAlongAxisGradKernel(const Context& dev_ctx,
   if (value_grad) {
     value_grad->Resize(index.dims());
     dev_ctx.template Alloc<T>(value_grad);
+#ifdef PADDLE_WITH_HIP
     auto* grad_data = value_grad->data<T>();
     int64_t grad_size = value_grad->numel();
-#ifdef PADDLE_WITH_HIP
     hipMemset(grad_data, 0, sizeof(T) * grad_size);
 #else
-    cudaMemset(grad_data, 0, sizeof(T) * grad_size);
+    // cudaMemset(grad_data, 0, sizeof(T) * grad_size);
+    phi::funcs::set_constant(dev_ctx, value_grad, 0);
 #endif
     if (reduce == "assign") {
       if (index_type == DataType::INT32) {