From 3bbff9e5ab964cf04897cebfc5e886a1113fef01 Mon Sep 17 00:00:00 2001 From: SangBin Cho Date: Thu, 21 Mar 2024 17:49:06 +0900 Subject: [PATCH] Fix 1D query issue from `_prune_hidden_states` (#3539) --- vllm/model_executor/layers/logits_processor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index baa113c3..e9d2a270 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -77,7 +77,6 @@ def _prune_hidden_states( hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata, ) -> torch.Tensor: - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) return hidden_states.index_select(0, sampling_metadata.selected_token_indices)