Skip to content

Commit de28b86

Browse files
committed
Optimize code format
Signed-off-by: skylee-01 <497627264@qq.com>
1 parent e7ad830 commit de28b86

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

vllm/model_executor/models/mlp_speculator.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -187,9 +187,11 @@ def generate_proposals(
187187
# TODO: not yet supporting top_k_tokens_per_head
188188
states = states.flatten(0, 1)
189189

190-
logits = self.logits_processor(
191-
self.head[head_index], states, sampling_metadata
192-
if self.sampling_metadata_is_required else None)
190+
if self.logits_processor:
191+
logits = self.logits_processor(self.head[head_index], states,
192+
sampling_metadata)
193+
else:
194+
logits = self.head[head_index](states)
193195

194196
output = self.sampler(logits, sampling_metadata)
195197
last_tokens = output.sampled_token_ids

0 commit comments

Comments
 (0)