[MXNET-1413] Adding Large Tensor support for sort operators

apache · Jun 20, 2019 · 496dd9f · 496dd9f
1 parent 2de0db0
commit 496dd9f
Show file tree

Hide file tree

Showing 4 changed files with 150 additions and 82 deletions.
diff --git a/src/operator/tensor/init_op.h b/src/operator/tensor/init_op.h
@@ -487,7 +487,7 @@ void EyeFill(const nnvm::NodeAttrs& attrs,
 
 struct range_fwd {
   template<typename DType>
-  MSHADOW_XINLINE static void Map(index_t i, int repeat, DType start, DType step,
+  MSHADOW_XINLINE static void Map(index_t i, index_t repeat, DType start, DType step,
                                   int req, DType* out) {
     KERNEL_ASSIGN(out[i], req, start + (i/repeat) * step);
   }