 
 template <typename DType>
 struct TempWorkspace {
-  index_t scores_temp_space;
+  size_t scores_temp_space;
   DType* scores;
-  index_t scratch_space;
+  size_t scratch_space;
   uint8_t* scratch;
-  index_t buffer_space;
+  size_t buffer_space;
   DType* buffer;
-  index_t nms_scratch_space;
+  size_t nms_scratch_space;
   uint32_t* nms_scratch;
-  index_t indices_temp_spaces;
+  size_t indices_temp_spaces;
   index_t* indices;
 };
 
-inline index_t ceil_div(index_t x, index_t y) {
+inline size_t ceil_div(size_t x, size_t y) {
   return (x + y - 1) / y;
 }
 
-inline index_t align(index_t x, index_t alignment) {
+inline size_t align(size_t x, size_t alignment) {
   return ceil_div(x, alignment)  * alignment;
 }
 

                 const int score_index,
                 Stream<gpu>* s) {
  const int n_threads = 512;
  const index_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) {

                      const index_t topk,
                      const int alignment,
                      TempWorkspace<DType>* workspace) {
  const size_t sort_scores_temp_space =
    mxnet::op::SortByKeyWorkspaceSize<DType, index_t, gpu>(num_elem, false, false);
  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);
