Update README.md

README.md

@@ -129,7 +129,13 @@ prompt = [
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-inputs = tokenizer.apply_chat_template(prompt, tokenize=True, add_generation_prompt=True, return_tensors="pt").cuda()
+inputs = tokenizer.apply_chat_template(
+  prompt,
+  tokenize=True,
+  add_generation_prompt=True,
+  return_tensors="pt",
+  return_dict=True,
+).to("cuda")
 
 model = AutoGPTQForCausalLM.from_pretrained(
   model_id,
@@ -138,7 +144,7 @@ model = AutoGPTQForCausalLM.from_pretrained(
   device_map="auto",
 )
 
-outputs = model.generate(inputs, do_sample=True, max_new_tokens=256)
+outputs = model.generate(**inputs, do_sample=True, max_new_tokens=256)
 print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
 ```