Empty cache after embedding to cpu (#9477)

2023-11-16 10:52:30 +08:00 · 2023-11-16 10:52:30 +08:00 · 731b0aaade
commit 731b0aaade
parent c487b53f21
1 changed files with 1 additions and 0 deletions
--- 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)