   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;
+  const size_t workspace_size = workspace.scores_temp_space +
+                                workspace.scratch_space +
+                                workspace.buffer_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);
+  Tensor<gpu, 1, double> scratch_memory = ctx.requested[box_nms_enum::kTempSpace]
+    .get_space_typed<gpu, 1, double>(mshadow::Shape1(ceil_div(workspace_size, sizeof(double))), s);
 
   // Populate workspace pointers
   workspace.scores = reinterpret_cast<DType*>(scratch_memory.dptr_);

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

