From 838452a33ffe031782864040006b2c74e00837fd Mon Sep 17 00:00:00 2001
From: Pei Zhang <zhangp1@ornl.gov>
Date: Fri, 20 Feb 2026 15:20:50 -0500
Subject: [PATCH] move norm to entry level in turbt only

---
 matey/models/turbt.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/matey/models/turbt.py b/matey/models/turbt.py
index 69ab547..8be1b85 100644
--- a/matey/models/turbt.py
+++ b/matey/models/turbt.py
@@ -249,15 +249,16 @@ def forward(self, data, state_labels, bcs, opts: ForwardOptionsBase):
             if imod<self.nhlevels-1:
                 x, blockdict=self.filterdata(x, blockdict=blockdict)
                 opts.blockdict = blockdict
+            else:
+                #self.debug_nan(x, message="input")
+                x, data_mean, data_std = normalize_spatiotemporal_persample(x)
+                #self.debug_nan(x, message="input after normalization")
             if imod>imod_bottom:
                 opts.imod -= 1
                 x_pred = self.forward(x, state_labels, bcs, opts)
             #x_input = x.clone()
             #T,B,C,D,H,W
             T, B, _, D, H, W = x.shape
-            #self.debug_nan(x, message="input")
-            x, data_mean, data_std = normalize_spatiotemporal_persample(x)
-            #self.debug_nan(x, message="input after normalization")
         ################################################################################
         if self.leadtime and leadtime is not None:
             leadtime = self.ltimeMLP[imod](leadtime)