Add batch dimension for cos, sin in unsloath test

Lightning-AI · Dec 31, 2024 · e20b7d9 · e20b7d9
1 parent 1b1b592
commit e20b7d9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_unsloth_executor.py b/tests/test_unsloth_executor.py
@@ -54,6 +54,8 @@ def test_unsloth_rope():
 
     B, nh, T, hs = 2, 32, 64, 16
     cos, sin = build_rope_cache(T, hs, device="cuda")
+    cos = cos.unsqueeze(0)
+    sin = sin.unsqueeze(0)
     q = torch.rand((B, nh, T, hs), device="cuda", requires_grad=True)
 
     def foo(x, cos, sin):