   int cur_mask = threadIdx.x + blockIdx.x * gridDim.x;
 
   int my_count = 0;
-  while(cur_mask < num_masks32)
+  while(cur_mask < (num_masks32 - 1))
   {
     my_count += __popc(masks32[cur_mask]);
 
     cur_mask += blockDim.x * gridDim.x;
   }
+
+  if(cur_mask == (num_masks32 - 1))
+  {
+    valid32_t last_mask = masks32[num_masks32 - 1];
+    int num_rows_last_mask = num_rows % BITS_PER_MASK32;
+    if(num_rows_last_mask == 0)
+      num_rows_last_mask = BITS_PER_MASK32;
+
+    for(int i = 0; i < num_rows_last_mask; ++i)
+    {
+      my_count += last_mask & gdf_valid_type(1);
+      last_mask >>= 1;
+    }
+  }
+
 
   int block_count = BlockReduce(temp_storage).Sum(my_count);
 

  return count;
}


__global__ 
void count_valid_bits(valid32_t const * const __restrict__ masks32, const int num_masks32, int * const __restrict__ global_count)
{
  typedef cub::BlockReduce<int, block_size> BlockReduce;
  __shared__ typename BlockReduce::TempStorage temp_storage;

  const size_t num_masks = gdf_get_num_chars_bitmask(num_rows);

  // Number of 4 byte types in the validity bit mask 
  size_t num_masks32 = std::ceil(static_cast<float>(num_masks) / RATIO);

  int h_count{0};
  if(num_masks32 > 0)
  {
    cudaStream_t count_stream;
    int * d_count;
    CUDA_TRY(cudaStreamCreate(&count_stream));
    // Cast validity buffer to 4 byte type
    valid32_t const * masks32 = reinterpret_cast<valid32_t const *>(masks);

    CUDA_TRY(cudaMalloc(&d_count, sizeof(int)));
    CUDA_TRY(cudaMemsetAsync(d_count, 0, sizeof(int),count_stream));

    const int grid_size = (num_masks32 + block_size - 1)/block_size;
    count_valid_bits<<<grid_size, block_size,0,count_stream>>>(masks32, num_masks32, num_rows, d_count);

    CUDA_TRY( cudaGetLastError() );

    CUDA_TRY(cudaMemcpyAsync(&h_count, d_count, sizeof(int), cudaMemcpyDeviceToHost,count_stream));
    CUDA_TRY(cudaStreamSynchronize(count_stream));
    CUDA_TRY(cudaStreamDestroy(count_stream));
  }

  // The final count of valid bits is the sum of the result from the
  // transform_reduce and the remainder masks
  *count = h_count;

  return GDF_SUCCESS;
}
