running mean and variance should only be updated during training

lucidrains · Nov 8, 2023 · 2bfcd1f · 2bfcd1f
1 parent 5584960
commit 2bfcd1f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/metnet3_pytorch/metnet3_pytorch.py b/metnet3_pytorch/metnet3_pytorch.py
@@ -799,7 +799,8 @@ def forward(
 
         # use a batchnorm to normalize each channel to mean zero and unit variance
 
-        _ = self.batchnorm_hrrr(hrrr_target)
+        if self.training:
+            _ = self.batchnorm_hrrr(hrrr_target)
 
         with freeze_batchnorm(self.batchnorm_hrrr) as frozen_batchnorm:
             normed_hrrr_pred = frozen_batchnorm(hrrr_pred)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'metnet3-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.5',
+  version = '0.0.6',
   license='MIT',
   description = 'MetNet 3 - Pytorch',
   author = 'Phil Wang',