huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-24 16:32:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Attempting to fix torch leak.

Browse Source
This commit is contained in:
Nicolas Patry 2023-08-12 09:06:49 +02:00
parent a2a913eec5
commit 89a4e723d2
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
server/text_generation_server/models/causal_lm.py
View File

@ -659,6 +659,7 @@ class CausalLM(Model):
# We finished all generations in the batch; there is no next batch # We finished all generations in the batch; there is no next batch
if stopped: if stopped:
torch.cuda.empty_cache()
return generations, None return generations, None
# Slice unused values from prefill # Slice unused values from prefill