Improve code snippet (#3)

- Improve code snippet (c4f91327b0b9f51cee01f41f9da8ed195e030378)
- Update README.md (c15b81007348c72c96b834df15b62bc57dc62965)
- Update README.md (d9b36fe104aaadadd78b092dc6f6b6e0263c3df4)

README.md

@@ -102,18 +102,15 @@ 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,
   low_cpu_mem_usage=True,
   device_map="auto",
-  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))
 ```