 };
 
 template <typename T>
-__global__ void MeanRunKernel(const T in_data, T* out_data, int N) {
+__global__ void MeanRunKernel(const T* in_data, T* out_data, int N) {
   int idx = blockDim.x * blockIdx.x + threadIdx.x;
+  T data = in_data[0];
   for (; idx < N; idx += blockDim.x * gridDim.x) {
-    out_data[idx] = in_data / (static_cast<T>(N));
+    out_data[idx] = data / (static_cast<T>(N));
   }
 }
 

    ops::MeanCUDAKernel<paddle::platform::CUDADeviceContext, double>,
    ops::MeanCUDAKernel<paddle::platform::CUDADeviceContext, plat::float16>);
REGISTER_OP_CUDA_KERNEL(
    mean_grad, ops::MeanGradKernel<paddle::platform::CUDADeviceContext, float>,
    ops::MeanGradKernel<paddle::platform::CUDADeviceContext, double>,
    ops::MeanGradKernel<paddle::platform::CUDADeviceContext, plat::float16>);

    auto IG = context.Output<Tensor>(framework::GradVarName("X"));
    IG->mutable_data<T>(context.GetPlace());

    auto in_data = OG->data<T>();
    auto size_prob = IG->numel();
    auto out_data = IG->data<T>();
    int threads = 512;
