   unsigned int j = blockIdx.y * blockDim.y + threadIdx.y;
   unsigned int k = blockIdx.z * blockDim.z + threadIdx.z;
 
-  if (i >= (c_Size.x - 1) || j >= (c_Size.y - 1)  || k >= (c_Size.z - 1))
+  if (i >= c_Size.x || j >= c_Size.y  || k >= c_Size.z)
       return;
 
   long int id   = (k     * c_Size.y + j)    * c_Size.x + i;

  unsigned int j = blockIdx.y * blockDim.y + threadIdx.y;
  unsigned int k = blockIdx.z * blockDim.z + threadIdx.z;

  if (i >= (c_Size.x - 1) || j >= (c_Size.y - 1)  || k >= (c_Size.z - 1))
      return;

  long int id   = (k     * c_Size.y + j)    * c_Size.x + i;

  multiply_by_beta_kernel <<< dimGrid, dimBlock >>> ( dev_in, dev_interm, beta);
  cudaDeviceSynchronize();
  CUDA_CHECK_ERROR;
//  std::cout << "Done multiplying" << std::endl;
  gradient_kernel <<< dimGrid, dimBlock >>> ( dev_interm, dev_grad_x, dev_grad_y, dev_grad_z);
  cudaDeviceSynchronize();
  CUDA_CHECK_ERROR;
//  std::cout << "Done computing gradient" << std::endl;
  magnitude_threshold_kernel <<< dimGrid, dimBlock >>> ( dev_grad_x, dev_grad_y, dev_grad_z, gamma);
  cudaDeviceSynchronize();
  CUDA_CHECK_ERROR;
//  std::cout << "Done projecting gradient onto L2 ball of radius gamma" << std::endl;

  // Rest of the iterations
  for (int iter=0; iter<NumberOfIterations; iter++)
