fix stella in main

predibase · Nov 8, 2024 · be29aea · be29aea
1 parent 414b4c9
commit be29aea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/server/lorax_server/models/flash_qwen2.py b/server/lorax_server/models/flash_qwen2.py
@@ -119,10 +119,10 @@ def is_row_parallel(self, layer_type: str) -> bool:
         return layer_type in ROW_PARALLEL
 
     def embed(self, batch) -> torch.Tensor:
-        adapter_meta = batch.adapter_meta
+        batch.prepare_for_prefill()
         prefill = False
         adapter_data = AdapterBatchData.from_meta(
-            adapter_meta, self.layer_to_adapter_weights, prefill, batch.prefill_head_indices
+            batch.adapter_meta, self.layer_to_adapter_weights, prefill, batch.prefill_head_indices
         )
         embedding, _ = self.forward(batch, adapter_data=adapter_data)
         return embedding.cpu().tolist()