hotfix mean, var assign

Qihoo360 · Sep 10, 2024 · 77fbd94 · 77fbd94
1 parent 4bc4e37
commit 77fbd94
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tensornet/layers/normalization_layer.py b/tensornet/layers/normalization_layer.py
@@ -120,9 +120,10 @@ def _increment_and_check_count():
                 _increment_and_check_count()
             else:
                 self.bn_statistics_push(False)
-        else:
-            mean = self.moving_mean
-            var = self.moving_variance
+            self.update_moments()
+
+        mean = self.moving_mean
+        var = self.moving_variance
 
         outputs = tf.nn.batch_normalization(x=inputs, mean=mean, variance=var, offset=self.beta, scale=self.gamma, variance_epsilon=self.epsilon)