Update modeling_grok2.py

modeling_grok2.py

@@ -104,10 +104,8 @@ class Grok2RMSNorm(nn.Module):
         self.eps = eps
 
     def forward(self, x):
-        orig = x.dtype
-        x = x.float()
-        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-        return (self.weight * x).to(orig)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        return self.weight * x * torch.rsqrt(variance + self.eps)
 
 
 # ── RoPE ──────────────────────────────────────────────────────────────────────