Update llama.py

unslothai · Jun 5, 2024 · 3771f5d · 3771f5d
1 parent 8380710
commit 3771f5d
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
@@ -1853,9 +1853,11 @@ def for_inference(model):
         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"):