diff --git a/paddle/fluid/operators/mean_op.cu b/paddle/fluid/operators/mean_op.cu index f0f895c08aac81252903ffd9a5ed1c931701adbb..cdd138d7bdc998fddef8af2e468e31b6f70aef49 100644 --- a/paddle/fluid/operators/mean_op.cu +++ b/paddle/fluid/operators/mean_op.cu @@ -31,10 +31,11 @@ struct DivideFunctor { }; template -__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(N)); + out_data[idx] = data / (static_cast(N)); } } @@ -85,7 +86,7 @@ class MeanCUDAGradKernel : public framework::OpKernel { auto IG = context.Output(framework::GradVarName("X")); IG->mutable_data(context.GetPlace()); - T in_data = OG[0]; + auto in_data = OG->data(); auto size_prob = IG->numel(); auto out_data = IG->data(); int threads = 512; @@ -105,6 +106,8 @@ REGISTER_OP_CUDA_KERNEL( ops::MeanCUDAKernel, ops::MeanCUDAKernel); REGISTER_OP_CUDA_KERNEL( - mean_grad, ops::MeanGradKernel, - ops::MeanGradKernel, - ops::MeanGradKernel); + mean_grad, + ops::MeanCUDAGradKernel, + ops::MeanCUDAGradKernel, + ops::MeanCUDAGradKernel);