From f2112b14adef00f2e679238e020fe31341de3b57 Mon Sep 17 00:00:00 2001
From: cm2435 <cm2435@bath.ac.uk>
Date: Tue, 6 Feb 2024 22:05:38 +0000
Subject: [PATCH] updated pre_patch for phi2

---
 unsloth/models/phi2.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/unsloth/models/phi2.py b/unsloth/models/phi2.py
index 8f8e0c12f3..0949b9bd19 100644
--- a/unsloth/models/phi2.py
+++ b/unsloth/models/phi2.py
@@ -169,8 +169,10 @@ class FastPhi2Model(FastLlamaModel):
     
     @staticmethod
     def pre_patch():
-        return 
-    
+        PhiAttention        .forward = Phi2Attention_fast_forward
+        PhiFlashAttention2  .forward = Phi2Attention_fast_forward
+    pass 
+
     @staticmethod
     def from_pretrained(
         model_name = "unsloth/llama-2-7b-bnb-4bit", #TODO: update me.