diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
index c1cc2b8282..b6e11b8bea 100644
--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -835,7 +835,7 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size):
             seq_dim = 0
 
         if reorder_strategy == ReorderStrategy.Striped:
-            seq_lens = shape[seq_dim]
+            seq_lens = tensor.shape[seq_dim]
             if seq_lens < (cp_size * stripe_size):
                 pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")