[MXNET-1413] Adding Large Tensor support for sort operators (#15170)

apache · Jun 21, 2019 · 8b5f376 · 8b5f376
1 parent 3f8fd00
commit 8b5f376
Show file tree

Hide file tree

Showing 4 changed files with 149 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);
   }