Fix GetReduceAces (#1460)

c3b462fc · Tatsuya Nishiyama · Tianqi Chen · fde98f0e · c3b462fc
Commit c3b462fc authored Jul 21, 2018 by Tatsuya Nishiyama Committed by Tianqi Chen Jul 20, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 1 additions and 1 deletions

nnvm/src/top/tensor/reduce.cc
+1 -1

No files found.
--- a/nnvm/src/top/tensor/reduce.cc
+++ b/nnvm/src/top/tensor/reduce.cc
@@ -48,7 +48,7 @@ inline TShape GetReduceAxes(const uint32_t indim,
  if (!exclude) return in_axis;
  TShape r_axis(indim - in_axis.ndim());
  for (unsigned i = 0, j = 0, k = 0; i < indim; ++i) {
-    if (i == in_axis[j]) {
+    if (j < in_axis.ndim() && i == in_axis[j]) {
        ++j;
        continue;
    }