Merge pull request #340 from EricLBuehler/fix_mistral_repeat

Fix mistral model repeat kv
EricLBuehler · May 22, 2024 · ea8dbec · ea8dbec
2 parents 6697d6c + f367820
commit ea8dbec
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mistralrs-core/src/models/mistral.rs b/mistralrs-core/src/models/mistral.rs
@@ -167,6 +167,7 @@ impl Attention {
         )?;
 
         let k = repeat_kv(k, self.num_kv_groups)?.contiguous()?;
+        let v = repeat_kv(v, self.num_kv_groups)?.contiguous()?;
 
         let mut attn_output = ScaledDotProductAttention.run_attention(
             &q,