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

fix(neuron): wrong assertion when batch_size==1

Browse Source
This commit is contained in:
David Corvoysier 2025-06-18 18:54:53 +00:00
parent bd1bdebb47
commit 82ebfd67bb
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
backends/neuron/server/text_generation_server/generator.py
View File

@ -341,7 +341,10 @@ class NeuronGenerator(Generator):
self.model = model self.model = model
if not isinstance(self.model, NeuronModelForCausalLM): if not isinstance(self.model, NeuronModelForCausalLM):
raise ValueError("The model must be a NeuronModelForCausalLM.") raise ValueError("The model must be a NeuronModelForCausalLM.")
if not model.neuron_config.continuous_batching: if (
model.neuron_config.batch_size > 1
and not model.neuron_config.continuous_batching
):
raise ValueError( raise ValueError(
"The neuron model must be compiled with continuous_batching=True." "The neuron model must be compiled with continuous_batching=True."
) )