fix: LayoutEnum import from cutlass.utils

2026-05-23 06:35:03 +00:00
parent bd968db5d4
commit 677245fc5c
1 changed files with 2 additions and 1 deletions
--- a/tests/unit/test_tmem_budget.py
+++ b/tests/unit/test_tmem_budget.py
@@ -6,7 +6,8 @@ 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
 from cutlass.cute.nvgpu import tcgen05
-from cutlass import BFloat16, Float32, LayoutEnum
+from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
+from cutlass.utils import LayoutEnum
 from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset

 def probe_hd(hd):