squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix 1D query issue from _prune_hidden_states (#3539)

Browse Source
This commit is contained in:
SangBin Cho 2024-03-21 17:49:06 +09:00 committed by GitHub
parent 6ebd02bdef
commit 3bbff9e5ab
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
vllm/model_executor/layers/logits_processor.py
View File

@ -77,7 +77,6 @@ def _prune_hidden_states(
hidden_states: torch.Tensor, hidden_states: torch.Tensor,
sampling_metadata: SamplingMetadata, sampling_metadata: SamplingMetadata,
) -> torch.Tensor: ) -> torch.Tensor:
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
return hidden_states.index_select(0, return hidden_states.index_select(0,
sampling_metadata.selected_token_indices) sampling_metadata.selected_token_indices)