                          ? static_cast<int>(ratio_w * out_img_idx + 0.5)
                          : static_cast<int>(ratio_w * out_img_idx);
 
-    T* in_pos;
-    if (data_layout == DataLayout::kNCHW) {
-      in_pos = &in[out_id_h * input_w + channel_id * in_img_size +
-                   in_img_idy * in_img_w + in_img_idx];
-    } else {
-      in_pos = &in[out_id_h * input_w + in_img_idy * in_img_w * num_channels +
-                   in_img_idx * num_channels + channel_id];
-    }
-    const T out_pos = out[out_id_h * output_w + out_id_w];
+    T* in_pos = &in[out_id_h * input_w + in_img_idy * in_img_w * num_channels +
+                    in_img_idx * num_channels + channel_id];
+
+    const T out_pos = out[tid];
     platform::CudaAtomicAdd(in_pos, out_pos);
   }
 }

      platform::GetGpuLaunchConfig1D(ctx.cuda_device_context(), pixelNum);

  if ("nearest" == interp_method) {
    KeNearestNeighborInterpBw<
        T><<<config.block_per_grid, config.thread_per_block, 0,
             ctx.cuda_device_context().stream()>>>(
        input_grad_data, in_h, in_w, n, in_chw, output_grad_data, out_h, out_w,
        n, out_chw, c, ratio_h, ratio_w, align_corners, data_layout);
  } else if ("bilinear" == interp_method) {
    const T align_type_value = (align_mode == 0 && !align_corners) ? 0.5f : 0;
    bool is_nchw = (data_layout == DataLayout::kNCHW) ? true : false;

}

template <typename T>
__global__ void KeNearestNeighborInterpNCHWBw(
    T* in, const size_t in_img_h, const size_t in_img_w, const T* out,
    const size_t out_img_h, const size_t out_img_w, const size_t nc,
    const float ratio_h, const float ratio_w, const bool align_corners) {
  int out_img_idx = threadIdx.x + blockIdx.x * blockDim.x;
  int out_img_idy = threadIdx.y + blockIdx.y * blockDim.y;
  int nc_id = threadIdx.z + blockIdx.z * blockDim.z;
  int nc_stride = blockDim.z * gridDim.z;

  // nearest_sampling by multiple read in_addr and write to out_addr
  int in_img_idx = (align_corners)
                       ? static_cast<int>(ratio_w * out_img_idx + 0.5)
                       : static_cast<int>(ratio_w * out_img_idx);
  int in_img_idy = (align_corners)
                       ? static_cast<int>(ratio_h * out_img_idy + 0.5)
                       : static_cast<int>(ratio_h * out_img_idy);

  int in_index = (nc_id * in_img_h + in_img_idy) * in_img_w + in_img_idx;
  int in_index_stride = nc_stride * in_img_h * in_img_w;

  int out_index = (nc_id * out_img_h + out_img_idy) * out_img_w + out_img_idx;
  int out_index_stride = nc_stride * out_img_h * out_img_w;

  // prevent from multiple threads writing
  if (out_img_idx < out_img_w && out_img_idy < out_img_h) {
    while (nc_id < nc) {
      T* in_pos = &in[in_index];
      const T out_pos = out[out_index];
      platform::CudaAtomicAdd(in_pos, out_pos);
      in_index += in_index_stride;
      out_index += out_index_stride;
      nc_id += nc_stride;
    }
  }
}

template <typename T>
__global__ void KeNearestNeighborInterpBw(
    T* in, const size_t in_img_h, const size_t in_img_w, const size_t input_h,
    const size_t input_w, const T* out, const size_t out_img_h,
    const size_t out_img_w, const size_t output_h, const size_t output_w,
    const size_t num_channels, const float ratio_h, const float ratio_w,
    const bool align_corners, FastDivModForInterpolate divmods) {
  int nthreads = output_h * output_w;
  int tid = blockIdx.x * blockDim.x + threadIdx.x;
  int stride = blockDim.x * gridDim.x;
  int in_img_size = in_img_h * in_img_w;
  int out_img_size = out_img_h * out_img_w;

  for (; tid < nthreads; tid += stride) {
    auto out_id_divmod = divmods.output_w_div.Divmod(tid);
    int out_id_h = out_id_divmod.val[0];
    int out_id_w = out_id_divmod.val[1];

    int channel_id = divmods.channels_div.Divmod(tid).val[1];
    auto outimg_id_divmod = divmods.output_wc_div.Divmod(out_id_w);
    int out_img_idy = outimg_id_divmod.val[0];
    int out_img_idx =
        divmods.channels_div.Divmod(outimg_id_divmod.val[1]).val[0];

    int in_img_idy = (align_corners)
                         ? static_cast<int>(ratio_h * out_img_idy + 0.5)
