From 37c53425c1f3a65f49e1cfa4e7cd80a67be277df Mon Sep 17 00:00:00 2001 From: mpikulski <206748156+ixlmar@users.noreply.github.com> Date: Fri, 13 Feb 2026 14:54:28 +0100 Subject: [PATCH] [TRTLLM-10030][chore] improve assert in sampler (#11475) Signed-off-by: ixlmar <206748156+ixlmar@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampling_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py index 66f04c3bef..0660ee7972 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py @@ -596,10 +596,10 @@ class SimpleGroupedStrategySampler(GroupedStrategySampler[Strategy]): beam_width_in = group_key[1] else: beam_width_in = 1 - if group_logit_indices is None: - assert logits.size(0) == beam_width_in * len(strategies) - else: + + if group_logit_indices is not None: logits = logits[group_logit_indices] + assert logits.size(0) == beam_width_in * len(strategies) assert all(strategy == group_key for strategy in strategies), "group must be consistent"