Update modeling_deepseek.py

modeling_deepseek.py

@@ -522,10 +522,14 @@ class DeepseekV3MoE(nn.Module):
         topk_idx, topk_weight = self.gate(hidden_states)
         hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
         flat_topk_idx = topk_idx.view(-1)
-        if not self.training:
-            y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
-        if self.config.n_shared_experts is not None:
-            y = y + self.shared_experts(identity)
+        # if not self.training:
+        y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)
+        try:
+            if self.config.n_shared_experts is not None:
+                y = y + self.shared_experts(identity)
+        except Exception as e:
+            if self.config.n_shared_experts is not None:
+                y = self.shared_experts(identity)
         return y
 
     @torch.no_grad()