Skip to content

GitLab

You need to sign in or sign up before continuing.
Commit a2cf1bdb authored by OlivierDehaene's avatar OlivierDehaene
Browse files

fix(server): empty_cache when stopped

parent c58a0c18
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
server/text_generation_server/models/flash_causal_lm.py
View file @ a2cf1bdb
...@@ -991,6 +991,7 @@ class FlashCausalLM(Model): ...@@ -991,6 +991,7 @@ class FlashCausalLM(Model):
if stopped: if stopped:
del batch del batch
torch.cuda.empty_cache()
# No need to return a batch if we know that all requests stopped # No need to return a batch if we know that all requests stopped
return generations, None return generations, None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment