diff --git a/rust/src/embeddings/local/qwen3.rs b/rust/src/embeddings/local/qwen3.rs index 5bb17e35..c903ee01 100644 --- a/rust/src/embeddings/local/qwen3.rs +++ b/rust/src/embeddings/local/qwen3.rs @@ -126,12 +126,20 @@ impl Qwen3Embed for Qwen3Embedder { for mini_text_batch in text_batch.chunks(batch_size) { let (token_ids, attention_mask) = tokenize_batch(&self.tokenizer, mini_text_batch, &self.device)?; + let embeddings: Tensor = { - let mut model = self.model.write().map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?; - model.forward(&token_ids, &attention_mask, 0)?.to_dtype(DType::F32)? + let mut model = self + .model + .write() + .map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?; + let result = model + .forward(&token_ids, &attention_mask, 0)? + .to_dtype(DType::F32)?; + model.clear_kv_cache(); + + result }; - self.model.write().map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?.clear_kv_cache(); let attention_mask = PooledOutputType::from(attention_mask); let attention_mask = Some(&attention_mask); let model_output = ModelOutput::Tensor(embeddings.clone());