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]