         }
 
 #pragma unroll
-        for (int ii = 0; ii < SZ; ii++) {
+        for (int ii = 0; ii < NumElementsPerThread; ii++) {
           weight[ii] += gradient[ii];
         }
 
 #pragma unroll
-        for (int ii = 0; ii < SZ; ii++) {
+        for (int ii = 0; ii < NumElementsPerThread; ii++) {
           int feature_dim = start_feature + ii * GPU_WARP_SIZE;
           if (feature_dim < stride) {
             grad_weight[weight_row + feature_dim] = static_cast<T>(weight[ii]);

  output[idx] = input[idx];
}

template <typename T, typename Tin>
__global__ void _GatherGradImpl(
    const Tin* input,
    const Tin* indices,

      }
      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]);
    }
  }
}

