Fix 1D query issue from _prune_hidden_states (vllm-project#3539)

stephanie-wang · Mar 21, 2024 · 3bbff9e · 3bbff9e
1 parent 6ebd02b
commit 3bbff9e
Showing 1 changed file with 0 additions and 1 deletion.
diff --git 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)