Improve code snippet

README.md

@@ -92,7 +92,7 @@ Alternatively, one may want to run that via `AutoAWQ` even though it's built on
 
 ```python
 import torch
-from autoawq import AutoAWQForCausalLM
+from awq import AutoAWQForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "hugging-quants/Meta-Llama-3.1-405B-Instruct-AWQ-INT4"
@@ -102,9 +102,6 @@ prompt = [
 ]
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-inputs = tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt").cuda()
-
 model = AutoAWQForCausalLM.from_pretrained(
   model_id,
   torch_dtype=torch.float16,
@@ -113,7 +110,8 @@ model = AutoAWQForCausalLM.from_pretrained(
   fuse_layers=True,
 )
 
-outputs = model.generate(inputs, do_sample=True, max_new_tokens=256)
+inputs = tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to('cuda')
+outputs = model.generate(**inputs, do_sample=True, max_new_tokens=256)
 print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
 ```