Repeat

EricLBuehler · Oct 24, 2024 · f34fdb9 · f34fdb9
1 parent b4dba1a
commit f34fdb9
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mistralrs-core/src/vision_models/mllama/text.rs b/mistralrs-core/src/vision_models/mllama/text.rs
@@ -429,7 +429,12 @@ impl MLlamaTextCrossAttention {
                 &k.contiguous()?.to_dtype(DType::F32)?,
                 &v.contiguous()?.to_dtype(DType::F32)?,
                 attention_mask
-                    .map(|m| m.to_dtype(DType::F32).unwrap())
+                    .map(|m| {
+                        m.to_dtype(DType::F32)
+                            .unwrap()
+                            .repeat((1, self.num_heads, 1, 1))
+                            .unwrap()
+                    })
                     .as_ref(),
                 None,
                 &self.sdpa_params,