Unverified Commit e4806d97 authored by duke's avatar duke Committed by GitHub
Browse files

[BugFix] Add embed_input_ids method to make QWenLMHeadModel a vllm model (#30674)


Signed-off-by: default avatarroot <iwzbi@zju.edu.cn>
Co-authored-by: default avatarroot <iwzbi@zju.edu.cn>
parent 4429d934
......@@ -281,6 +281,9 @@ class QWenBaseModel(nn.Module):
self.transformer.make_empty_intermediate_tensors
)
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.transformer.wte(input_ids)
def compute_logits(
self,
hidden_states: torch.Tensor,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment