diff --git a/tests/test_wo_a.py b/tests/test_wo_a.py
index 17f81a34..2fb21062 100644
--- a/tests/test_wo_a.py
+++ b/tests/test_wo_a.py
@@ -47,11 +47,11 @@ def test_inverse_rope():
     rope_dim = ROPE_DIM
     half_rope = rope_dim // 2
     base = 10000.0
-    inv_freq = 1.0 / (base ** (torch.arange(0, half_rope, dtype=torch.float32) / half_rope))
+    inv_freq = 1.0 / (base ** (torch.arange(0, half_rope, dtype=torch.float32, device=DEVICE) / half_rope))
 
-    pos = torch.arange(max_pos, dtype=torch.float32)
+    pos = torch.arange(max_pos, dtype=torch.float32, device=DEVICE)
     freqs = torch.outer(pos, inv_freq)  # (max_pos, half_rope)
-    cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1)  # (max_pos, rope_dim)
+    cos_sin_cache = torch.cat([freqs.cos(), freqs.sin(), ], dim=-1)  # (max_pos, rope_dim)
 
     # Random attention output
     o = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE) * 2.0