Update modeling_grok2.py

modeling_grok2.py

@@ -180,7 +180,8 @@ class Grok2Expert(nn.Module):
     def forward(self, x):
         device = self.w1.weight.device
         x = x.to(device)
-        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+        h = F.silu(self.w1(x)) * self.w3(x.to(self.w3.weight.device))
+        return self.w2(h.to(self.w2.weight.device))
 
 
 # ── Sparse MoE ────────────────────────────────────────────────────────────────