Fix GetReduceAces (#1460)

apache · Jul 20, 2018 · d32b509 · d32b509
1 parent 5c84a98
commit d32b509
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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;
     }