fix

smallcloudai · Dec 29, 2023 · d4e31d1 · d4e31d1
1 parent 7abe19a
commit d4e31d1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/self_hosting_machinery/finetune/modelling/flash_sa.py b/self_hosting_machinery/finetune/modelling/flash_sa.py
@@ -164,7 +164,7 @@ def _forward(
         k = einops.rearrange(k, "b t (h d) -> b h t d", h=self.num_key_value_heads)
         v = einops.rearrange(v, "b t (h d) -> b t h d", h=self.num_key_value_heads)
 
-        cos, sin = self.rotary_emb(q, seq_len=k.shape[-2])
+        cos, sin = self.rotary_emb(v, seq_len=k.shape[-2])
         q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
 
         q = einops.rearrange(q, "b h t d -> b t h d")

diff --git a/self_hosting_machinery/finetune/scripts/aux/model.py b/self_hosting_machinery/finetune/scripts/aux/model.py
@@ -154,7 +154,8 @@ def forward(
             input,
             return_dict=False,
             output_attentions=False,
-            output_hidden_states=False
+            output_hidden_states=False,
+            use_cache=False
         )[0]
         return logits