Skip to content

Commit 91f3f3f

Browse files
committed
fix: huggingface backend does not properly pad inputs
1 parent 03017e7 commit 91f3f3f

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

mellea/backends/huggingface.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,9 @@ def _generate_from_raw(
447447
prompts = [self.formatter.print(action) for action in actions]
448448

449449
# batch-encoding call is deprecated in favor of this
450-
inputs = self._tokenizer(prompts, return_tensors="pt").to(self._device)
450+
inputs = self._tokenizer(prompts, return_tensors="pt", padding=True).to(
451+
self._device
452+
)
451453

452454
if format is None:
453455
outputs = self._model.generate( # type: ignore

0 commit comments

Comments
 (0)