Update modeling_moe_mistral.py

modeling_moe_mistral.py

@@ -196,9 +196,7 @@ class FeedForward(nn.Module):
         )
 
     def forward(self, x):
-        device = x.device
-        x = x.to(self.w1.weight.device)
-        return self.w2(F.silu(self.w1(x)) * self.w3(x)).to(device)
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
 
 
 class MoE(nn.Module):
@@ -217,8 +215,9 @@ class MoE(nn.Module):
         orig_shape = x.shape
         x = x.view(-1, x.shape[-1])
 
-        scores = self.gate(x).softmax(dim=-1)
+        scores = self.gate(x)
         expert_weights, expert_indices = torch.topk(scores, self.num_experts_per_token, dim=-1)
+        expert_weights = expert_weights.softmax(dim=-1)
         flat_expert_indices = expert_indices.view(-1)
 
         x = x.repeat_interleave(self.num_experts_per_token, dim=0)