Update README.md

README.md

@@ -15,12 +15,12 @@ This model was trained using AutoTrain. For more information, please visit [Auto
 
 ```python
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import MistralConfig, AutoTokenizer
 
 model_path = "PATH_TO_THIS_REPO"
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(
+model = MistralConfig.from_pretrained(
     model_path,
     device_map="auto",
     torch_dtype='auto'