mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-05-24 09:51:19 +08:00
[Bugfix] Fix isinstance check for tensor types in _load_prompt_embeds to use dtype comparison (#21612)
Signed-off-by: Alexandre Juan <a.juan@netheos.net>
This commit is contained in:
parent
a55c95096b
commit
2f6e6b33fb
@ -957,9 +957,11 @@ class OpenAIServing:
|
|||||||
def _load_and_validate_embed(embed: bytes) -> EmbedsPrompt:
|
def _load_and_validate_embed(embed: bytes) -> EmbedsPrompt:
|
||||||
tensor = torch.load(io.BytesIO(base64.b64decode(embed)),
|
tensor = torch.load(io.BytesIO(base64.b64decode(embed)),
|
||||||
weights_only=True)
|
weights_only=True)
|
||||||
assert isinstance(
|
assert isinstance(tensor, torch.Tensor) and tensor.dtype in (
|
||||||
tensor,
|
torch.float32,
|
||||||
(torch.FloatTensor, torch.BFloat16Tensor, torch.HalfTensor))
|
torch.bfloat16,
|
||||||
|
torch.float16,
|
||||||
|
)
|
||||||
if tensor.dim() > 2:
|
if tensor.dim() > 2:
|
||||||
tensor = tensor.squeeze(0)
|
tensor = tensor.squeeze(0)
|
||||||
assert tensor.dim() == 2
|
assert tensor.dim() == 2
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user