diff --git a/tensorrt_edgellm/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/tensorrt_edgellm/models/qwen3_5_moe/modeling_qwen3_5_moe.py index e4875033..79d79204 100644 --- a/tensorrt_edgellm/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/tensorrt_edgellm/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -76,11 +76,6 @@ def __init__(self, config: ModelConfig, layer_idx: int) -> None: 1, module_name=f"{prefix}.shared_expert_gate") - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - routed = super().forward(hidden_states) - shared = self.shared_expert(hidden_states) - shared_gate = torch.sigmoid(self.shared_expert_gate(hidden_states)) - return routed + shared * shared_gate class Qwen3_5MoeDecoderLayer(Qwen3_5DecoderLayer):