zwischenraum commited on
Commit
120bf03
·
verified ·
1 Parent(s): 7081022

Update README.md

Browse files

For me the code did not work out of the box. The issue is that `apply_chat_template` with `tokenize=True` + `return_tensors="pt"` returns a `BatchEncoding` object, but `model.generate()` expects a plain tensor. Small fix should do the trick.

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,