Fix DSV3 training (#324)

2025-07-25 09:47:06 -07:00
parent 08e4dd2fc5
commit 64763adeeb
1 changed files with 3 additions and 1 deletions
@@ -303,7 +303,9 @@ def group_expert_select(
    group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
    k = n_group - topk_group
    group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-    scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
+    scores = mx.put_along_axis(
+        scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+    )
    scores = mx.flatten(scores, -2, -1)

    k = top_k