Fix wrong batch size of thinking_mask (#4296)

Co-authored-by: K11OntheBoat <“ruianmaidanglao@163.com”>
Co-authored-by: xiegegege <46314656+xiegegege@users.noreply.github.com>
This commit is contained in:
K11OntheBoat
2025-09-28 14:56:42 +08:00
committed by GitHub
parent 3cef851468
commit 7b6cb72ab2

View File

@@ -196,7 +196,7 @@ def post_process_normal(
"""Post-processing steps after completing a single token generation."""
# handle vl:
if model_output.think_end_id != -1:
thinking_mask = model_output.enable_thinking
thinking_mask = model_output.enable_thinking[: sampler_output.sampled_token_ids.shape[0]]
exists_think_end = (sampler_output.sampled_token_ids == model_output.think_end_id) & thinking_mask
paddle.assign(
paddle.where(