diff --git a/mlx_lm/models/deepseek_v3.py b/mlx_lm/models/deepseek_v3.py
index 264abd1..2ae57a4 100644
--- a/mlx_lm/models/deepseek_v3.py
+++ b/mlx_lm/models/deepseek_v3.py
@@ -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