ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Empty cache after embedding to cpu (#9477)

Browse source
This commit is contained in:
Yuwen Hu 2023-11-16 10:52:30 +08:00 committed by GitHub
parent c487b53f21
commit 731b0aaade
1 changed files with 1 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
python/llm/src/bigdl/llm/transformers/embedding.py
View file

@ -23,4 +23,5 @@ class LLMEmbedding(torch.nn.Embedding):
def forward(self, x: Tensor): def forward(self, x: Tensor):
if self.weight.device != 'cpu': if self.weight.device != 'cpu':
self.to('cpu') self.to('cpu')
torch.xpu.empty_cache()
return super().forward(x.to('cpu')).to(x.device) return super().forward(x.to('cpu')).to(x.device)