Fix 1D query issue from _prune_hidden_states (#3539)

2024-03-21 17:49:06 +09:00
parent 6ebd02bdef
commit 3bbff9e5ab
1 changed files with 0 additions and 1 deletions
--- 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)