Update README.md

README.md

@@ -51,7 +51,7 @@ Use the code below to get started with the model.
 
 ### Load model directly
 ```python
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 
 base_model_id = "johnsnowlabs/JSL-MedLlama-3-8B-v2.0"