Fix token_indices: repeat each token ID top_k times, not arange

2026-05-17 21:22:11 +00:00
parent a478ca4746
commit 1acf01fc1a
1 changed files with 1 additions and 1 deletions
--- 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]