kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix: reshape token_ids for lp in torch backend (#4239)

Browse Source 
reshape token_ids

Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
This commit is contained in:
Erin 2025-05-12 17:43:47 -07:00 committed by GitHub
parent 035d915fea
commit 4becf32360
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorrt_llm/_torch/pyexecutor/model_engine.py
View File

@ -2035,6 +2035,7 @@ class PyTorchModelEngine(ModelEngine):
# Reshape to align w/ the shape used in the TRT backend,
# so the same logit processors can be used across both backends.
logits_row = logits_row.view(1, 1, -1)
token_ids = [token_ids]
for lp in logits_processors:
lp_params = inspect.signature(lp).parameters