refine elementwise

7 years ago · 5e04b64fa0
parent 5ac72d9501
commit 5e04b64fa0
1 changed files with 6 additions and 4 deletions
--- a/paddle/operators/elementwise_op_function.h
+++ b/paddle/operators/elementwise_op_function.h
@ -103,10 +103,12 @@ class MidWiseTransformIterator<T, platform::CPUDeviceContext> {

  MidWiseTransformIterator<T, platform::CPUDeviceContext>& operator++() {
    ++j_;
-    i_ = j_ / post_;
-    if (UNLIKELY(i_ == n_)) {
-      j_ = 0;
-      i_ = 0;
+    if (UNLIKELY(j_ == post_)) {
+      ++i_;
+      if (UNLIKELY(i_ == n_)) {
+        j_ = 0;
+        i_ = 0;
+      }
    }
    return *this;
  }