Update modeling_grok2.py

modeling_grok2.py

@@ -178,6 +178,8 @@ class Grok2Expert(nn.Module):
         self.w3 = nn.Linear(hidden_size, moe_intermediate_size, bias=False)
 
     def forward(self, x):
+        device = self.w1.weight.device
+        x = x.to(device)
         return self.w2(F.silu(self.w1(x)) * self.w3(x))
 
 
@@ -345,4 +347,4 @@ class Grok1ForCausalLM(PreTrainedModel, GenerationMixin):
 
 # ── Register ──────────────────────────────────────────────────────────────────
 AutoConfig.register("grok2", Grok2Config)
-AutoModelForCausalLM.register(Grok2Config, Grok1ForCausalLM)
+AutoModelForCausalLM.register(Grok2Config, Grok1ForCausalLM)