Only update BN moving stats during training

nhynes · Oct 19, 2018 · 0f8c0dc · 0f8c0dc
1 parent bdd15ea
commit 0f8c0dc
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/nnvm/src/top/nn/nn.cc b/nnvm/src/top/nn/nn.cc
@@ -358,17 +358,17 @@ axis to be the last item in the input shape.
                                   inputs[3], mean, param.momentum);
       new_var = MakeMomentumNode(n->attrs.name + "_var_mom",
                                  inputs[4], var_unbiased, param.momentum);
+
+      new_mean = MakeNode("_assign", n->attrs.name + "_mean_update",
+                          {inputs[3], new_mean});
+      new_var = MakeNode("_assign", n->attrs.name + "_var_update",
+                         {inputs[4], new_var});
     } else {
       mean = inputs[3];
       var = inputs[4];
       new_mean = inputs[3];
       new_var = inputs[4];
     }
-
-    new_mean = MakeNode("_assign", n->attrs.name + "_mean_update",
-                        {inputs[3], new_mean});
-    new_var = MakeNode("_assign", n->attrs.name + "_var_update",
-                       {inputs[4], new_var});
     mean = compiler::ExpandBiasToMatchAxis(mean, in_dim, 1, ax);
     var = compiler::ExpandBiasToMatchAxis(var, in_dim, 1, ax);