From 731b0aaade615a223014dd188224a2f1ec4a2e43 Mon Sep 17 00:00:00 2001
From: Yuwen Hu <54161268+Oscilloscope98@users.noreply.github.com>
Date: Thu, 16 Nov 2023 10:52:30 +0800
Subject: [PATCH] Empty cache after embedding to cpu (#9477)

---
 python/llm/src/bigdl/llm/transformers/embedding.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/llm/src/bigdl/llm/transformers/embedding.py b/python/llm/src/bigdl/llm/transformers/embedding.py
index 38aa4db7..aa99e2d7 100644
--- a/python/llm/src/bigdl/llm/transformers/embedding.py
+++ b/python/llm/src/bigdl/llm/transformers/embedding.py
@@ -23,4 +23,5 @@ class LLMEmbedding(torch.nn.Embedding):
     def forward(self, x: Tensor):
         if self.weight.device != 'cpu':
             self.to('cpu')
+            torch.xpu.empty_cache()
         return super().forward(x.to('cpu')).to(x.device)