From 1acf01fc1a472ee4952a6d8bffe97bde1771559e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 21:22:11 +0000 Subject: [PATCH] Fix token_indices: repeat each token ID top_k times, not arange --- tests/test_pipeline_real_weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py index 71f27ede..3ee66e87 100644 --- a/tests/test_pipeline_real_weights.py +++ b/tests/test_pipeline_real_weights.py @@ -70,7 +70,7 @@ def main(): flat_ids = safe_ids.reshape(-1) flat_weights = safe_weights.reshape(-1) num_slots = NUM_TOKENS * TOP_K - token_indices = torch.arange(num_slots, device=DEVICE) + token_indices = torch.arange(NUM_TOKENS, device=DEVICE).unsqueeze(1).expand(-1, TOP_K).reshape(-1) sort_idx = flat_ids.argsort(stable=True) sorted_ids = flat_ids[sort_idx]