From d2fa16aafa35d002801c0dc2ee842b67c97e6e86 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Wed, 5 Jun 2024 20:57:08 +1000
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 7dec8624e1..9c9b45dbae 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1853,9 +1853,11 @@ class FastLlamaModel:
         pass
 
         # Wrap model.generate
-        model._unwrapped_old_generate = model.generate
-        model.generate = _wrap_fast_inference(model.generate, device_type, dtype, model)
-
+        if model.generate.__name__ != "_fast_generate":
+            model._unwrapped_old_generate = model.generate
+            model.generate = _wrap_fast_inference(model.generate, device_type, dtype, model)
+        pass
+        
         # Patch tokenizer to pad to the left
         internal_model = model
         while hasattr(internal_model, "model"):