[rust] Make tensor contiguous in rotary embedding (#3436)

deepjavalibrary · Aug 21, 2024 · 5c418be · 5c418be
1 parent 99ed366
commit 5c418be
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/extensions/tokenizers/rust/src/models/gte.rs b/extensions/tokenizers/rust/src/models/gte.rs
@@ -134,8 +134,8 @@ impl RotaryEmbedding {
         let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?;
         let cos = self.cos.narrow(0, 0, seq_len)?;
         let sin = self.sin.narrow(0, 0, seq_len)?;
-        let q_embed = rotary_emb::rope(q, &cos, &sin)?;
-        let k_embed = rotary_emb::rope(k, &cos, &sin)?;
+        let q_embed = rotary_emb::rope(&q.contiguous()?, &cos, &sin)?;
+        let k_embed = rotary_emb::rope(&k.contiguous()?, &cos, &sin)?;
         Ok((q_embed, k_embed))
     }
 }