Update

vllm-project · Oct 4, 2024 · 4a05af3 · 4a05af3
1 parent 8e423d8
commit 4a05af3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/vllm/model_executor/models/gemma2_embedding.py b/vllm/model_executor/models/gemma2_embedding.py
@@ -8,7 +8,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.gemma2 import Gemma2Model
 from vllm.model_executor.pooling_metadata import PoolingMetadata
-from vllm.sequence import PoolerOutput
+from vllm.sequence import IntermediateTensors, PoolerOutput
 
 
 class Gemma2EmbeddingModel(nn.Module):
@@ -36,10 +36,12 @@ def forward(
         positions: torch.Tensor,
         kv_caches: List[torch.Tensor],
         attn_metadata: AttentionMetadata,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         return self.model.forward(input_ids, positions, kv_caches,
-                                  attn_metadata, inputs_embeds)
+                                  attn_metadata, intermediate_tensors,
+                                  inputs_embeds)
 
     def pooler(
         self,