diff --git a/tests/unit/test_fmha_v3_stage_c_full.py b/tests/unit/test_fmha_v3_stage_c_full.py index 6026fb9f..d1d356b9 100644 --- a/tests/unit/test_fmha_v3_stage_c_full.py +++ b/tests/unit/test_fmha_v3_stage_c_full.py @@ -253,7 +253,7 @@ class FmhaV3StageC: frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) for j in range(frg_cnt): - for k in cutlass.range(cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True): + for k in range(cute.size(tTMEM_LOADrS_frg, mode=[0])): row_max = cute.arch.fmax(row_max, tTMEM_LOADrS_frg[k, j] * scale_log2) row_max_safe = row_max