From a5ccea7403a58f2a5549eae47a9f2746d6766f41 Mon Sep 17 00:00:00 2001 From: Weimin Wang <301118019+weiminwang-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 18:28:13 +0800 Subject: [PATCH] fix mypy error in sampler_strategy.py Signed-off-by: Weimin Wang <301118019+weiminwang-nv@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index fc04afbb51e4..51b69c1d4c24 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -351,7 +351,7 @@ def sample( logits, beam_width_in=cast(int, beam_width_in), beam_width_out=cast(int, beam_width_out), - row_stride=cast(int, row_stride), + row_stride=row_stride, beam_search_args=group_metadata, temperature=cast(float, temperature), early_stopping=cast(int, early_stopping), @@ -1341,7 +1341,7 @@ def sample_grouped_strategies( # variable-beam-width step; the op slices down to the live beams. rows_per_request = beam_width_in if strategies and strategies[0][0] == "beam_search": - rows_per_request = cast(BeamSearch, strategies[0]).row_stride + rows_per_request = strategies[0].row_stride assert logits.size(0) == rows_per_request * len(strategies) else: assert group_logit_indices.size(0) == beam_width_in * len(strategies)