diff --git a/tests/test_e2e_decode_b200.py b/tests/test_e2e_decode_b200.py index 48e8a2c3..79d96e68 100644 --- a/tests/test_e2e_decode_b200.py +++ b/tests/test_e2e_decode_b200.py @@ -274,8 +274,8 @@ def main(): runners = {} weights = {} - # Test with 3 representative layers: C128A, C4A, SWA - test_layers = [0, 2, 60] + # Test with all 61 layers (shared experts only) + test_layers = list(range(NUM_LAYERS)) for layer_id in test_layers: p = f"model.layers.{layer_id}"