 
   int64_t pre_idx = idx / (stride * size);
   int64_t dim_idx = idx % (stride * size) / stride;
-  int64_t begin_idx = idx + (delta * pre_idx - dim_idx) * stride;
-
+  IndexT src_dim_idx = index[dim_idx];
+  int64_t input_idx = idx + (delta * pre_idx + src_dim_idx - dim_idx) * stride;
+  paddle::platform::CudaAtomicAdd(&input_grad[input_idx], output_grad[idx]);
+}
+
+template <typename T>
+__global__ void index_select_grad_init(T* input_grad, int64_t N) {
+  int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= N) {
+    return;
+  }
   input_grad[idx] = 0.0;
-  for (int64_t i = 0; i < nums; i++) {
-    if (index[i] == dim_idx) {
-      input_grad[idx] += output_grad[begin_idx + i * stride];
-    }
-  }
 }
 
 template <typename DeviceContext, typename T>

    dim = dim >= 0 ? dim : dim + input_dim.size();
    auto stride_dim = framework::stride(input_dim);
    int64_t stride = stride_dim[dim];
    int64_t size = input_dim[dim];
    int64_t delta = output_dim[dim] - size;

    const auto& index_type = index->type();
    bool index_type_match = index_type == framework::proto::VarType::INT64 ||

    index_select,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, float>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, double>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext,
                               paddle::platform::float16>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, int>,
    ops::IndexSelectCUDAKernel<paddle::platform::CUDADeviceContext, int64_t>);
REGISTER_OP_CUDA_KERNEL(
    index_select_grad,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext, float>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext, double>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext,
                                   paddle::platform::float16>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext, int>,
    ops::IndexSelectGradCUDAKernel<paddle::platform::CUDADeviceContext,
                                   int64_t>);
