   CUDA_KERNEL_ASSERT(start <= N);
   CUDA_KERNEL_ASSERT(end <= N);
 
-  for (int i = threadIdx.x; i < post; i += blockDim.x) {
-    for (int line = start; line < end; ++line) {
-      grad_out[line * post + i] = grad_in[group * post + i];
+  if (ExactBlock) {
+    grad_out += threadIdx.x;
+    grad_in += threadIdx.x;
+
+    for (int line = start + threadIdx.y; line < end; line += blockDim.y) {
+      grad_out[line * post] = grad_in[group * post];
+    }
+  } else {
+    for (int i = threadIdx.x; i < post; i += blockDim.x) {
+      for (int line = start; line < end; ++line) {
+        grad_out[line * post + i] = grad_in[group * post + i];
+      }
     }
   }
 }

                len_length,
                dataToReduceSize);
      } else {
        sparse_length_sum_kernel<T, IndexType>
            <<<len_length, 32, 0, context_.cuda_stream()>>>(
                in_data,
                out_data,
                prefix_sum_length_data,

    int N = output_0dim;
    int post = segmentGradsInput.size_from_dim(1);

    auto maxThreads =
        GetDeviceProperty(CaffeCudaGetDevice()).maxThreadsPerBlock;

    if (post <= maxThreads) {
      int multiple = std::min(maxThreads / post, 16);
      dim3 block(post, multiple);

      length_sum_gradient_kernel<T, true>
          <<<len_length, block, 0, context_.cuda_stream()>>>(

              in_data, out_data, prefix_sum_length_data, N, post, len_length);
    } else {
      length_sum_gradient_kernel<T, false>
          <<<len_length, maxThreads, 0, context_.cuda_stream()>>>(
              in_data, out_data, prefix_sum_length_data, N, post, len_length);
    }

