diff --git a/single_shot_inference.py b/single_shot_inference.py
index fc8f553d..886b74db 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -343,13 +343,13 @@ class HcHead:
         # Flatten: (T, n_hc * H)
         X_flat = X_L.reshape(T, self.K).bfloat16()
 
-        # Unweighted RMSNorm (same as in mHC)
+        # Unweighted RMSNorm on flattened residual (FP32 for numerical stability)
         X_f = X_flat.float()
         rms_inv = X_f.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
-        X_normed = (X_f * rms_inv).bfloat16()
+        X_normed = X_f * rms_inv  # Keep FP32 for the linear
 
-        # Linear projection: (T, K) @ (4, K).T → (T, 4)
-        mixes = torch.nn.functional.linear(X_normed, self.fn.bfloat16()).float()
+        # Linear projection: (T, K) @ (4, K).T → (T, 4) in FP32
+        mixes = torch.nn.functional.linear(X_normed, self.fn.float()).float()
 
         # Apply scale + bias + sigmoid + eps
         pre = torch.sigmoid(mixes * self.scale + self.base.float().unsqueeze(0)) + self.eps