From 222714c495832a8909f61f422ab667eb3502e981 Mon Sep 17 00:00:00 2001 From: Gabriel de Marmiesse Date: Tue, 13 Jan 2026 18:07:02 +0100 Subject: [PATCH] Fix memory leak by using no_grad in all threads (#16) `torch.no_grad` is thread local, I now run the autoregressive generation in a separate thread, in addition of the thread running mimi, and forgot to add `torch.no_grad` to the function inside this thread. --- pocket_tts/models/tts_model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 7ef44e9..027f456 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -261,6 +261,7 @@ class TTSModel(nn.Module): conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight) return conditioning + @torch.no_grad def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue): """Worker thread function for decoding audio latents from queue with immediate streaming.""" try: @@ -509,6 +510,7 @@ class TTSModel(nn.Module): generation_thread = threading.Thread(target=run_generation, daemon=True) generation_thread.start() + @torch.no_grad def _autoregressive_generation( self, model_state: dict, max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue ):