Update README.md

README.md

@@ -107,13 +107,14 @@ pip3 install git+https://github.com/casper-hansen/AutoAWQ.git@1c5ccc791fa2cb0697
 ```python
 from awq import AutoAWQForCausalLM
 from transformers import AutoTokenizer
+import torch
 
 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}
@@ -154,7 +155,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'])