diff --git a/tests/unit/test_d3_swa_mask.py b/tests/unit/test_d3_swa_mask.py
index c066d289..df4ad64f 100644
--- a/tests/unit/test_d3_swa_mask.py
+++ b/tests/unit/test_d3_swa_mask.py
@@ -58,7 +58,7 @@ def _run_fmha(q_3d, k_3d, v, m, s_k, hd, use_smem_p=False):
         c_tile.zero_(); lse_tensor.zero_()
         mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
         mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
-        mLSE = ct.from_dlset(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
+        mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
         compiled(mQ, mK, mV, mC, stream, mLSE)
         o_unnorm[:, pv*pv_n_tile:(pv+1)*pv_n_tile] = c_tile[:,:,0].float()