remove component that causes low tpu duty cycle on multi-host

AI-Hypercomputer · Jul 18, 2024 · 658647a · 658647a
1 parent 53eb544
commit 658647a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
@@ -686,7 +686,7 @@ def _generate_thread(self, idx: int):
         decode_state = generate_engine.insert(
             new_request.prefill_result, decode_state, slot=slot
         )
-        delete_pytree(new_request.prefill_result)
+        del new_request.prefill_result
         new_request.generate_timestep_added = generate_timestep
         new_request.complete = np.zeros(
             (generate_engine.samples_per_slot,), dtype=np.bool_