Update modeling_densebackward_olmoe0125.py

modeling_densebackward_olmoe0125.py

@@ -43,8 +43,8 @@ class DenseBackwardOlmoeSparseMoeBlock(OlmoeSparseMoeBlock):
         
         # 3) build hard & ste masks
         mask_hard = F.one_hot(selected_experts, num_classes=self.num_experts).sum(dim=1).to(dtype)     # (N, num_experts)
-        mask_ste = mask_hard + (routing_weights - routing_weights.detach())
-        # mask_ste = mask_hard + (router_logits - router_logits.detach())
+        #mask_ste = mask_hard + (routing_weights - routing_weights.detach())
+        mask_ste = mask_hard + (router_logits - router_logits.detach())
 
         # 4) compute gated weights = π * mask, then optionally renormalize
         gated = routing_weights * mask_ste                              # zero-out non-TopK