normalization of hrrr pred and targets should use same statistics

lucidrains · Nov 8, 2023 · 84c336a · 84c336a
1 parent c533fcc
commit 84c336a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/metnet3_pytorch/metnet3_pytorch.py b/metnet3_pytorch/metnet3_pytorch.py
@@ -799,11 +799,11 @@ def forward(
 
         # use a batchnorm to normalize each channel to mean zero and unit variance
 
-        normed_hrrr_target = self.batchnorm_hrrr(hrrr_target)
-
         with freeze_batchnorm(self.batchnorm_hrrr) as frozen_batchnorm:
             normed_hrrr_pred = frozen_batchnorm(hrrr_pred)
 
+        normed_hrrr_target = self.batchnorm_hrrr(hrrr_target)
+
         # proposed loss gradient rescaler from section 4.3.2
 
         normed_hrrr_pred = self.mse_loss_scaler(normed_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.3',
+  version = '0.0.4',
   license='MIT',
   description = 'MetNet 3 - Pytorch',
   author = 'Phil Wang',