From bfa46c38d5724b7064ff2b3365efcfa406281a63 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Wed, 10 Jun 2020 11:44:28 +0800
Subject: [PATCH] bn supports reverse_space, test=develop (#24988)

---
 paddle/fluid/pybind/op_function_generator.cc | 3 +++
 python/paddle/fluid/dygraph/nn.py            | 3 ++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/pybind/op_function_generator.cc b/paddle/fluid/pybind/op_function_generator.cc
index fc60b6302e7..89ebb925363 100644
--- a/paddle/fluid/pybind/op_function_generator.cc
+++ b/paddle/fluid/pybind/op_function_generator.cc
@@ -52,6 +52,9 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
 std::map<std::string, std::set<std::string>> op_outs_map = {
     {"fake_quantize_dequantize_moving_average_abs_max",
      {"Out", "OutScale", "OutAccum", "OutState"}},
+    {"batch_norm",
+     {"Y", "MeanOut", "VarianceOut", "SavedMean", "SavedVariance",
+      "ReserveSpace"}},
 };
 
 // NOTE(zhiqiu): Commonly, the outputs in auto-generated OP function are
diff --git a/python/paddle/fluid/dygraph/nn.py b/python/paddle/fluid/dygraph/nn.py
index 7cb4702fedc..680e47bf422 100644
--- a/python/paddle/fluid/dygraph/nn.py
+++ b/python/paddle/fluid/dygraph/nn.py
@@ -1310,9 +1310,10 @@ class BatchNorm(layers.Layer):
                      self._fuse_with_relu, "use_global_stats",
                      self._use_global_stats, 'trainable_statistics',
                      self._trainable_statistics)
-            batch_norm_out, _, _, _, _ = core.ops.batch_norm(
+            batch_norm_out, _, _, _, _, _ = core.ops.batch_norm(
                 input, self.weight, self.bias, self._mean, self._variance,
                 mean_out, variance_out, *attrs)
+
             return dygraph_utils._append_activation_in_dygraph(
                 batch_norm_out, act=self._act)
 
-- 
GitLab