fix import path for tcgen05

2026-05-23 06:34:30 +00:00
parent afec55ad72
commit bd968db5d4
1 changed files with 2 additions and 1 deletions
--- a/tests/unit/test_tmem_budget.py
+++ b/tests/unit/test_tmem_budget.py
@@ -4,7 +4,8 @@ Prints find_tmem_tensor_col_offset(tOtO) and related shapes so we can
 plan the SMEM-P path and verify TMEM fits in 512 columns at hd=512.
 """
 import torch, math
-import cutlass, cutlass.cute as cute, cutlass.utils as utils, cutlass.nvgpu.tcgen05 as tcgen05
+import cutlass, cutlass.cute as cute, cutlass.utils as utils
+from cutlass.cute.nvgpu import tcgen05
 from cutlass import BFloat16, Float32, LayoutEnum
 from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset