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