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',