Update README.md

README.md

@@ -105,13 +105,17 @@ pip3 install git+https://github.com/casper-hansen/AutoAWQ.git@1c5ccc791fa2cb0697
 ```python
 from awq import AutoAWQForCausalLM
 from transformers import AutoTokenizer
+import torch
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
 
 model_name_or_path = "TheBloke/Mistral-7B-v0.1-AWQ"
 
 # Load model
 model = AutoAWQForCausalLM.from_quantized(model_name_or_path, fuse_layers=True,
-                                          trust_remote_code=False, safetensors=True)
-tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=False)
+                                          trust_remote_code=False, safetensors=True, device=device)
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=False, device=device)
 
 prompt = "Tell me about AI"
 prompt_template=f'''{prompt}
@@ -152,7 +156,8 @@ pipe = pipeline(
     temperature=0.7,
     top_p=0.95,
     top_k=40,
-    repetition_penalty=1.1
+    repetition_penalty=1.1,
+    device=device
 )
 
 print(pipe(prompt_template)[0]['generated_text'])