Update modeling_grok2.py

modeling_grok2.py

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