fix

Signed-off-by: Austin Liu <[email protected]>
linkedin · Jan 16, 2025 · 83dcd59 · 83dcd59
1 parent 460cfaf
commit 83dcd59
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/transformers/test_rope.py b/test/transformers/test_rope.py
@@ -58,7 +58,7 @@ def test_correctness(
     atol,
     rtol,
 ):
-    rotary_emb = LlamaRotaryEmbedding(LlamaConfig(head_dim=head_dim), device=device)
+    rotary_emb = LlamaRotaryEmbedding(config=LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), device=device)
 
     _tensor_q = torch.randn((bsz, seq_len, num_q_heads, head_dim), device=device).transpose(1, 2).to(dtype)
 
@@ -134,7 +134,7 @@ def test_functional_correctness(
     k1 = _k.clone().requires_grad_(True)
     k2 = _k.clone().requires_grad_(True)
 
-    rotary_emb = LlamaRotaryEmbedding(LlamaConfig(head_dim=head_dim), device=device)
+    rotary_emb = LlamaRotaryEmbedding(config=LlamaConfig(num_kv_heads=num_kv_heads, head_dim=head_dim), device=device)
 
     pos_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0)
     if expand_position_ids: