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):