From 0e41d117a0afe269e50be1b2b0d847182bc1ab9e Mon Sep 17 00:00:00 2001 From: ducviet00 Date: Tue, 13 May 2025 20:53:39 +0700 Subject: [PATCH] fix: correct batchsize for tasks in enc_dec models Signed-off-by: ducviet00 --- tensorrt_llm/models/enc_dec/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/models/enc_dec/model.py b/tensorrt_llm/models/enc_dec/model.py index 338f16c54a..7ffe875bed 100644 --- a/tensorrt_llm/models/enc_dec/model.py +++ b/tensorrt_llm/models/enc_dec/model.py @@ -912,7 +912,7 @@ def prepare_inputs(self, dtype=trt.int32, shape=[-1, 1], dim_range=OrderedDict([ - ('batch_size', bs_range), + ('batch_size', [bs_range]), ('broadcast_dim', [1]), ])) prompt_vocab_size = Tensor(name='prompt_vocab_size',