diff --git a/tests/unit/test_d1_lse_verify.py b/tests/unit/test_d1_lse_verify.py
index 01c2a6d1..90200ff1 100644
--- a/tests/unit/test_d1_lse_verify.py
+++ b/tests/unit/test_d1_lse_verify.py
@@ -59,13 +59,14 @@ def test_lse(hd=64):
     torch.cuda.synchronize()
 
     kernel_lse = lse_tensor.cpu()
+    ref_lse_cpu = ref_lse.cpu()
 
     # Compare
-    lse_err = (kernel_lse - ref_lse).abs()
+    lse_err = (kernel_lse - ref_lse_cpu).abs()
     print(f'  LSE max err: {lse_err.max().item():.6f}')
     print(f'  LSE mean err: {lse_err.mean().item():.6f}')
     print(f'  Kernel LSE[:8]: {kernel_lse[:8].tolist()}')
-    print(f'  Ref LSE[:8]:    {ref_lse[:8].tolist()}')
+    print(f'  Ref LSE[:8]:    {ref_lse_cpu[:8].tolist()}')
 
     # Check O too
     ref_unnorm = attn_exp @ v.float()