Update README.md

README.md

@@ -127,10 +127,12 @@ pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory
 Then you can load the model using
 
 ```python
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("HiTZ/GoLLIE-7B")
-model = AutoModelForCausalLM.from_pretrained("HiTZ/GoLLIE-7B", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("HiTZ/GoLLIE-7B", trust_remote_code=True, torch_dtype=torch.bfloat16)
+model.to("cuda")
 ```
 
 Read our [🚀 Example Jupyter Notebooks](notebooks/) to learn how to easily define guidelines, generate model inputs and parse the output!