Only narrow when needed + deactivate the kv cache.

1 files changed, 6 insertions, 2 deletions

diff --git a/candle-core/examples/llama/main.rs b/candle-core/examples/llama/main.rs
index 9d70921c..5a8a15d3 100644
--- a/candle-core/examples/llama/main.rs
+++ b/candle-core/examples/llama/main.rs
@@ -24,7 +24,7 @@ mod var_store;
 mod weights;
 
 const CONTEXT_SIZE: usize = 512;
-const USE_KV_CACHE: bool = true;
+const USE_KV_CACHE: bool = false;
 const START_PROMPT: &str = r"
 EDWARD:
 I wonder how our princely father 'scaped,
@@ -268,7 +268,11 @@ impl CausalSelfAttention {
 
     fn apply_rotary_emb(&self, x: &Tensor, freqs_cis: &Tensor) -> Result<Tensor> {
         let mut dims = x.dims().to_vec();
-        let freqs_cis = freqs_cis.narrow(1, freqs_cis.dims()[1] - dims[1], dims[1])?;
+        let freqs_cis = if dims[1] < CONTEXT_SIZE {
+            freqs_cis.narrow(1, CONTEXT_SIZE - dims[1], dims[1])?
+        } else {
+            freqs_cis.clone()
+        };
         let v = dims.pop().unwrap();
         dims.push(v / 2);
         dims.push(2);