Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: better warmup error

This commit is contained in:
OlivierDehaene 2023-10-25 10:18:58 +02:00
parent f9910d13e2
commit 96a982ad8f
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/models/flash_causal_lm.py
View File

@ -670,7 +670,7 @@ class FlashCausalLM(Model):
self.device, self.device,
) )
_, batch = self.generate_token(batch) _, batch = self.generate_token(batch)
except Exception as e: except torch.cuda.OutOfMemoryError as e:
raise RuntimeError( raise RuntimeError(
f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. " f"Not enough memory to handle {len(batch.input_ids)} prefill tokens. "
f"You need to decrease `--max-batch-prefill-tokens`" f"You need to decrease `--max-batch-prefill-tokens`"