Fix T scope in post_block

2026-05-31 00:02:29 +00:00
parent a493f72681
commit 3ecfbcba57
1 changed files with 2 additions and 1 deletions
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -136,7 +136,8 @@ class mHCBlock:
        post_term = post_mix.unsqueeze(-1) * F_out.unsqueeze(1).float()
        residual_next = mixed_residual + post_term
        # Emergency: per-residual RMSNorm (routed MoE missing)
-        rn_f = residual_next.reshape(T, self.n_hc, -1)
+        _T = residual_next.shape[0]
+        rn_f = residual_next.reshape(_T, self.n_hc, -1)
        rms = rn_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
        return (rn_f * rms).bfloat16()
 # RoPE