Skip to content

Commit e78fc4b

Browse files
hmellorAkshat-Tripathi
authored andcommitted
Fix failing MyGemma2Embedding test (vllm-project#13820)
Signed-off-by: Harry Mellor <[email protected]>
1 parent fecd1c2 commit e78fc4b

File tree

1 file changed

+1
-6
lines changed

1 file changed

+1
-6
lines changed

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
# SPDX-License-Identifier: Apache-2.0
22

3-
from typing import Iterable, List, Optional, Tuple, Union
3+
from typing import Iterable, Optional, Tuple, Union
44

55
import torch
66
import torch.nn as nn
77

8-
from vllm.attention import AttentionMetadata
98
from vllm.config import VllmConfig
109
from vllm.model_executor.layers.pooler import Pooler, PoolingType
1110
from vllm.model_executor.models.gemma2 import Gemma2Model
@@ -37,16 +36,12 @@ def forward(
3736
self,
3837
input_ids: torch.Tensor,
3938
positions: torch.Tensor,
40-
kv_caches: List[torch.Tensor],
41-
attn_metadata: AttentionMetadata,
4239
intermediate_tensors: Optional[IntermediateTensors] = None,
4340
inputs_embeds: Optional[torch.Tensor] = None,
4441
) -> Union[torch.Tensor, IntermediateTensors]:
4542
hidden_states = self.model(
4643
input_ids,
4744
positions,
48-
kv_caches,
49-
attn_metadata,
5045
intermediate_tensors=intermediate_tensors,
5146
inputs_embeds=inputs_embeds,
5247
)

0 commit comments

Comments
 (0)