From 4becf32360f69aaf324e321e06db5190e2babe33 Mon Sep 17 00:00:00 2001
From: Erin <14718778+hchings@users.noreply.github.com>
Date: Mon, 12 May 2025 17:43:47 -0700
Subject: [PATCH] fix: reshape token_ids for lp in torch backend (#4239)

reshape token_ids

Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 6aa98ed027..7db587a11b 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -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