       num_indices));
 
   dim3 block(GPU_WARP_SIZE, 4);
-  dim3 grid(CeilDiv(num_indices, 4), CeilDiv(stride, 128));
-
-  hipLaunchKernelGGL(_GatherGradImpl, dim3(grid), dim3(block), 0, 0, 
+  dim3 grid(CeilDiv(num_indices, 4), CeilDiv(stride, GridDim::maxElementsPerThread * GPU_WARP_SIZE));
+  if (param_itrs == 1)
+  {
+    hipLaunchKernelGGL(HIP_KERNEL_NAME(_GatherAxis0GradImpl<T, Tin, GridDim::maxElementsPerThread>), dim3(grid), dim3(block), 0, 0,
       indices_data_sorted.get(),
       original_indices_sorted.get(),
       grad_data,
       output_data,
       num_indices,
       num_inputs,
-      param_itrs,
-      stride);
+      stride); 
+  } else {
+    hipLaunchKernelGGL(HIP_KERNEL_NAME(_GatherGradImpl<T, Tin, GridDim::maxElementsPerThread>), dim3(grid), dim3(block), 0, 0,
+        indices_data_sorted.get(),
+        original_indices_sorted.get(),
+        grad_data,
+        output_data,
+        num_indices,
+        num_inputs,
+        param_itrs,
+        stride);
+  }
 }
 
 #define SPECIALIZED_GRAD_IMPL2(T)           \

    int64_t stride) {
  int idx = blockIdx.x * 4 + threadIdx.y;

  const int SZ = 4;
  if (idx < numel && (idx == 0 || input[idx] != input[idx - 1])) {
    do {
      for (int itr = 0; itr < param_itrs; ++itr) {
        const int start_feature = threadIdx.x + blockIdx.y * blockDim.x * SZ;
        const int weight_row = itr * input_numel + ((int)input[idx]) * stride;  //the offset of the input
        const int grad_row = (itr * numel + ((int)indices[idx])) * stride;      //the offset of the gradient

        float gradient[SZ];
        float weight[SZ];

#pragma unroll
        for (int ii = 0; ii < SZ; ii++) {
          int feature_dim = start_feature + ii * GPU_WARP_SIZE;
          if (feature_dim < stride) {
            gradient[ii] = static_cast<float>(grad_output[grad_row + feature_dim]);

      }
      idx++;
    } while (idx < numel && input[idx] == input[idx - 1]);
  }
}

// Special optimization for the case which the gather is on axis=0
template <typename T, typename Tin, int NumElementsPerThread>
__global__ void _GatherAxis0GradImpl(
    const Tin* input,
    const Tin* indices,
    const T* grad_output,
    T* grad_weight,
    int64_t numel,
    int64_t input_numel,
    int64_t stride)
{
  int idx = blockIdx.x * 4 + threadIdx.y;

  if (idx < numel && (idx == 0 || input[idx] != input[idx - 1])) {
    const int start_feature = threadIdx.x + blockIdx.y * blockDim.x * NumElementsPerThread;
    const int weight_row = ((int)input[idx]) * stride;  //the offset of the input

    float weight[NumElementsPerThread];
    for (int ii = 0; ii < NumElementsPerThread; ii++) {
      int feature_dim = start_feature + ii * GPU_WARP_SIZE/4;
      if (feature_dim < stride)
        weight[ii] = static_cast<float>(grad_weight[weight_row + feature_dim]);
    }

    do {
        const int grad_row = ((int)indices[idx]) * stride;      //the offset of the gradient
        float gradient[NumElementsPerThread];

#pragma unroll
        for (int ii = 0; ii < NumElementsPerThread; ii++) {
          int feature_dim = start_feature + ii * GPU_WARP_SIZE/4;
          if (feature_dim < stride) {
            gradient[ii] = static_cast<float>(grad_output[grad_row + feature_dim]);
            weight[ii] += gradient[ii];
          }
        }
        idx++;
    } while (idx < numel && input[idx] == input[idx - 1]);
    
#pragma unroll
    for (int ii = 0; ii < NumElementsPerThread; ii++) {
      int feature_dim = start_feature + ii * GPU_WARP_SIZE/4;
      if (feature_dim < stride)
        grad_weight[weight_row + feature_dim] = static_cast<T>(weight[ii]);
    }
  }
}

