[fix] Fix for neuron rolling batch CI error (deepjavalibrary#1938)

tosterberg · May 21, 2024 · 4d283a8 · 4d283a8
1 parent 28d0480
commit 4d283a8
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/transformers_neuronx_scheduler/optimum_neuron_scheduler.py b/engines/python/setup/djl_python/transformers_neuronx_scheduler/optimum_neuron_scheduler.py
@@ -162,8 +162,6 @@ def make_generations(
                 generated_tokens=slot_request.slot.generated_tokens,
                 finish_reason=finish_reason,
                 seed=slot_request.slot.seed)
-            # mark the slot as available
-            slot_request.slot.clear()
         generation = Generation(
             request_id=slot_request.slot.request_id,
             prefill_tokens=None,
@@ -174,6 +172,9 @@ def make_generations(
                               in slot_request.slot.special_tokens),
             generated_text=generated_text,
             speculated_generations=SpeculatedGenerationsQueue())
+        if finish_reason is not None:
+            # mark the slot as available
+            slot_request.slot.clear()
         return generation, finish_reason
 
     def _generate_token(self, inputs: GenerationInputs) -> List[Generation]: