           N * D, P->data<float>(), P->mutable_data<float>(), &context_);
     } else {
       ProbCrossEntropyKernel<<<
-          CAFFE_GET_BLOCKS(N),
+          std::min(N, CAFFE_MAXIMUM_NUM_BLOCKS),
           CAFFE_CUDA_NUM_THREADS,
           0,
           context_.cuda_stream()>>>(

    float* dXdata,
    const float* weights) {
  if (weights == NULL) {
    CUDA_1D_KERNEL_LOOP(i, N) {
      for (int j = 0; j < D; j++) {
        int idx = i * D + j;
        dXdata[idx] = Pdata[idx] - labeldata[idx];
      }
    }
  } else {
    CUDA_1D_KERNEL_LOOP(i, N) {
      float weight = weights[i];
      for (int d = 0; d < D; d++) {
        int idx = i * D + d;
        dXdata[idx] = (Pdata[idx] - labeldata[idx]) * weight;
      }
    }
  }
}

    const float* labeldata,
    const float* weights,
    float* Ydata) {
  typedef cub::BlockReduce<float, CAFFE_CUDA_NUM_THREADS> BlockReduce;
  __shared__ typename BlockReduce::TempStorage temp_storage;

  for (int i = blockIdx.x; i < N; i += gridDim.x) {
    float weight = weights ? weights[i] : 1.0;
    float sum = 0.0;
    float total_prob = 0.0;
    for (int j = threadIdx.x; j < D; j += blockDim.x) {
      int idx = i * D + j;
      CUDA_KERNEL_ASSERT(labeldata[idx] >= 0);
      total_prob += labeldata[idx];
      sum += -logf(max(Pdata[idx], FLT_MIN)) * labeldata[idx] * weight;
    }
    float tot = BlockReduce(temp_storage).Sum(sum);
    __syncthreads();
    float total_prob_sum = BlockReduce(temp_storage).Sum(total_prob);
    if (threadIdx.x == 0) {
      Ydata[i] = tot;
      // Sanity check
      CUDA_KERNEL_ASSERT(abs(1.0 - total_prob_sum) < 1e-5f);
    }
    __syncthreads();
  }
}

