Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -155,7 +155,7 @@ input_ids = tokenizer.apply_chat_template(
155
  add_generation_prompt=True,
156
  return_tensors="pt",
157
  tokenize=True,
158
- ).to(model.device)
159
 
160
  output = model.generate(
161
  input_ids,
 
155
  add_generation_prompt=True,
156
  return_tensors="pt",
157
  tokenize=True,
158
+ )["input_ids"].to(model.device)
159
 
160
  output = model.generate(
161
  input_ids,