Merge pull request #11 from mj-will/fix-batch-norm-memory-leak

Fix memory leak in batch norm
bayesiains · Jul 11, 2020 · 6244a7d · 6244a7d
2 parents 84dc029 + 36dcabc
commit 6244a7d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nflows/transforms/normalization.py b/nflows/transforms/normalization.py
@@ -103,8 +103,8 @@ def forward(self, inputs, context=None):
 
         if self.training:
             mean, var = inputs.mean(0), inputs.var(0)
-            self.running_mean.mul_(1 - self.momentum).add_(mean * self.momentum)
-            self.running_var.mul_(1 - self.momentum).add_(var * self.momentum)
+            self.running_mean.mul_(1 - self.momentum).add_(mean.detach() * self.momentum)
+            self.running_var.mul_(1 - self.momentum).add_(var.detach() * self.momentum)
         else:
             mean, var = self.running_mean, self.running_var