From 677245fc5cae5c8a116d8eb2f400e157efa4c7df Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 23 May 2026 06:35:03 +0000
Subject: [PATCH] fix: LayoutEnum import from cutlass.utils

---
 tests/unit/test_tmem_budget.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_tmem_budget.py b/tests/unit/test_tmem_budget.py
index 99859011..fd1b08e5 100644
--- 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):