address reviews

apache · Jul 15, 2019 · c5e5363 · c5e5363
1 parent c8bee52
commit c5e5363
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/src/operator/nn/softmax-inl.h b/src/operator/nn/softmax-inl.h
@@ -75,7 +75,7 @@ inline void Softmax(Stream<cpu> *s, DType *in, OType *out,
   index_t sa = stride[axis];
 
   #pragma omp parallel for
-  for (int i = 0; i < static_cast<int>(N); ++i) {
+  for (index_t i = 0; i < N; ++i) {
     index_t base = unravel_dot(i, sshape, stride);
 
     DType mmax = negate ? -in[base] : in[base];
@@ -125,8 +125,8 @@ inline void SoftmaxWithLength(Stream<cpu> *s, DType *in, OType *out, IType *leng
   index_t sa = stride[axis];
 
   #pragma omp parallel for
-  for (int i = 0; i < static_cast<int>(N); ++i) {
-    IType len = length[i];
+  for (index_t i = 0; i < N; ++i) {
+    index_t len = static_cast<index_t>(length[i]);
     index_t base = unravel_dot(i, sshape, stride);
 
     DType mmax = negate ? -in[base] : in[base];
@@ -135,7 +135,7 @@ inline void SoftmaxWithLength(Stream<cpu> *s, DType *in, OType *out, IType *leng
       val = negate ? -in[base + j*sa] : in[base + j*sa];
       if (mmax < val) mmax = val;
     }
-    for (int j = len; j < M; ++j) {
+    for (index_t j = len; j < M; ++j) {
       out[base + j*sa] = OType(0.0f);
     }
 
@@ -202,7 +202,7 @@ inline void SoftmaxGrad(Stream<cpu> *s, OType *out, OType *ograd,
   index_t sa = stride[axis];
 
   #pragma omp parallel for
-  for (int i = 0; i < static_cast<int>(N); ++i) {
+  for (index_t i = 0; i < N; ++i) {
     index_t base = unravel_dot(i, sshape, stride);
 
     AType sum = AType(0);
@@ -244,9 +244,9 @@ inline void SoftmaxWithLengthGrad(Stream<cpu> *s, OType *out, OType *ograd,
   index_t sa = stride[axis];
 
   #pragma omp parallel for
-  for (int i = 0; i < static_cast<int>(N); ++i) {
+  for (index_t i = 0; i < N; ++i) {
     index_t base = unravel_dot(i, sshape, stride);
-    IType len = length[i];
+    index_t len = static_cast<index_t>(length[i]);
 
     AType sum = AType(0);
     for (index_t j = 0; j < len; ++j) {

diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
@@ -7871,7 +7871,6 @@ def get_output_names_callback(name, arr):
         except mx.base.MXNetError:
             # skip errors since test is to check all names
             pass
-        print(output_names)
         for output_name, expected_name in zip(output_names, expected_names):
             assert output_name == expected_name