From 9714bf74382d026be0da9f3a2ad2b93b2ef35797 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Tue, 12 Dec 2023 13:07:38 +1100
Subject: [PATCH] Update llama.py

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

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index e4b5118e62..1f0046ec06 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -614,7 +614,7 @@ class FastLlamaModel:
         rope_scaling = None,
     ):
         SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported()
-        print_unsloth_message("Mistral")
+        print_unsloth_message("Llama")
         FastLlamaModel.pre_patch()
 
         if dtype is None:
@@ -734,6 +734,7 @@ class FastLlamaModel:
         use_gradient_checkpointing = True,
         random_state = 3407,
         max_seq_length = 2048,
+        **kwargs,
     ):
         assert(max_seq_length <= model.max_seq_length)
 
@@ -759,6 +760,7 @@ class FastLlamaModel:
             bias           = "none",
             task_type      = TaskType.CAUSAL_LM,
             layers_to_transform = layers_to_transform,
+            **kwargs,
         )
 
         model = prepare_model_for_kbit_training(