From 6c8618c75443a08c4ec0304cc54acfae74b2ddcd Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jul 2024 12:08:49 -0700
Subject: [PATCH] Update mistral.py

---
 unsloth/models/mistral.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index 6eb3fccfab..d7376d9521 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -78,6 +78,9 @@ def MistralAttention_fast_forward(
     if past_key_value is not None:
         kv_seq_len += past_key_value[0].shape[-2]
 
+    # Extend RoPE dynamically to fit in VRAM
+    self.rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len)
+
     if position_ids is None:
         cos = self.rotary_emb.cos_cached
         sin = self.rotary_emb.sin_cached