Update utils.py

utils.py

@@ -22,6 +22,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import datasets
 from datasets import load_dataset
 import evaluate
+from transformers import LlamaForCausalLM, LlamaTokenizer
 
 
 
@@ -91,7 +92,8 @@ def load_tokenizer_and_model_Blaize(base_model, load_8bit=True):
     
     tokenizer = LlamaTokenizer.from_pretrained(base_model, add_eos_token=True, use_auth_token=True)
     model = LlamaForCausalLM.from_pretrained(base_model, load_in_8bit=True, device_map="auto")
-    #model.eval()
+    model = prepare_model_for_int8_training(model)
+
     return tokenizer,model, device