From 52e17bf50c2e89e3a29136e4679d312d9c1aa995 Mon Sep 17 00:00:00 2001
From: guosheng <guosheng@baidu.com>
Date: Mon, 29 Jan 2018 13:56:36 +0800
Subject: [PATCH] Fix the unit test of weight normalization

---
 paddle/operators/reduce_op.cc                            | 9 ++++++---
 .../paddle/v2/fluid/tests/test_weight_normalization.py   | 5 ++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/paddle/operators/reduce_op.cc b/paddle/operators/reduce_op.cc
index 4a06babeda..84f24a9095 100644
--- a/paddle/operators/reduce_op.cc
+++ b/paddle/operators/reduce_op.cc
@@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/operators/reduce_op.h"
-#include "paddle/operators/net_op.h"
 
 namespace paddle {
 namespace operators {
@@ -38,10 +37,14 @@ class ReduceOp : public framework::OperatorWithKernel {
         dim, x_rank,
         "The dim should be in the range [-rank(input), rank(input)).");
     bool reduce_all = ctx->Attrs().Get<bool>("reduce_all");
+    bool keep_dim = ctx->Attrs().Get<bool>("keep_dim");
     if (reduce_all) {
-      ctx->SetOutputDim("Out", {1});
+      if (keep_dim)
+        ctx->SetOutputDim(
+            "Out", framework::make_ddim(std::vector<int64_t>(x_rank, 1)));
+      else
+        ctx->SetOutputDim("Out", {1});
     } else {
-      bool keep_dim = ctx->Attrs().Get<bool>("keep_dim");
       auto dims_vector = vectorize(x_dims);
       if (keep_dim || x_rank == 1) {
         dims_vector[dim] = 1;
diff --git a/python/paddle/v2/fluid/tests/test_weight_normalization.py b/python/paddle/v2/fluid/tests/test_weight_normalization.py
index 1ac54ea305..80ad8285d8 100644
--- a/python/paddle/v2/fluid/tests/test_weight_normalization.py
+++ b/python/paddle/v2/fluid/tests/test_weight_normalization.py
@@ -112,9 +112,8 @@ class TestWeightNormalization(unittest.TestCase):
             [
                 self.assertTrue(
                     numpy.allclose(
-                        numpy.array(actual_output), expect_output, atol=0.001))
-                for expect_output, actual_output in zip(expect_output,
-                                                        actual_output)
+                        numpy.array(actual), expect, atol=0.001))
+                for expect, actual in zip(expect_output, actual_output)
             ]
 
 
-- 
GitLab