                       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);

  workspace.nms_scratch_space = align(NMS<DType>::THRESHOLD / (sizeof(uint32_t) * 8) *
                                      num_batch * topk * sizeof(uint32_t), alignment);

  const index_t workspace_size = workspace.scores_temp_space +
                                 workspace.scratch_space +
                                 workspace.nms_scratch_space +
                                 workspace.indices_temp_spaces;

  // Obtain the memory for workspace
  Tensor<gpu, 1, uint8_t> scratch_memory = ctx.requested[box_nms_enum::kTempSpace]
    .get_space_typed<gpu, 1, uint8_t>(mshadow::Shape1(workspace_size), s);

  // Populate workspace pointers
  workspace.scores = reinterpret_cast<DType*>(scratch_memory.dptr_);


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

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

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

