                       const index_t topk,
                       const int alignment,
                       TempWorkspace<DType>* workspace) {
-  const index_t sort_scores_temp_space =
+  const size_t sort_scores_temp_space =
     mxnet::op::SortByKeyWorkspaceSize<DType, index_t, gpu>(num_elem, false, false);
-  const index_t sort_topk_scores_temp_space =
+  const size_t sort_topk_scores_temp_space =
     mxnet::op::SortByKeyWorkspaceSize<DType, index_t, gpu>(topk, false, false);
   workspace->scratch_space = align(std::max(sort_scores_temp_space, sort_topk_scores_temp_space),
                                    alignment);


template <typename DType>
struct TempWorkspace {
  index_t scores_temp_space;
  DType* scores;
  index_t scratch_space;
  uint8_t* scratch;
  index_t buffer_space;
  DType* buffer;
  index_t nms_scratch_space;
  uint32_t* nms_scratch;
  index_t indices_temp_spaces;
  index_t* indices;
};

inline index_t ceil_div(index_t x, index_t y) {
  return (x + y - 1) / y;
}

inline index_t align(index_t x, index_t alignment) {
  return ceil_div(x, alignment)  * alignment;
}


                 const int score_index,
                 Stream<gpu>* s) {
  const int n_threads = 512;
  const size_t max_blocks = 320;
  index_t N = source.shape_.Size();
  const auto blocks = std::min(ceil_div(N, n_threads), max_blocks);
  if (topk > 0) {
