Update README.md

README.md

@@ -17,13 +17,13 @@ language:
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 model_path = "Rakuten/RakutenAI-7B-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="auto")
 model = PeftModel.from_pretrained(
     model,
-    "Kendamarron/math-problem-generator_RakutenAI-7B-instruct_lora",
-    torch_dtype=torch.bfloat16
+    "Kendamarron/math-problem-generator_RakutenAI-7B-instruct_lora"
 )
 model.eval()