fix(gpt): stream mode dim mismatch (fix #606) (#607)

2noise · Jul 20, 2024 · 51ec0c7 · 51ec0c7
1 parent a0e6cd8
commit 51ec0c7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ChatTTS/model/gpt.py b/ChatTTS/model/gpt.py
@@ -416,6 +416,7 @@ def generate(
         )
         inputs_ids_buf.narrow(1, 0, progress).copy_(inputs_ids)
         del inputs_ids
+        inputs_ids = inputs_ids_buf.narrow(1, 0, progress)
 
         pbar: Optional[tqdm] = None
 
@@ -430,8 +431,6 @@ def generate(
 
         for i in range(max_new_token):
 
-            inputs_ids = inputs_ids_buf.narrow(1, 0, progress)
-
             model_input = self._prepare_generation_inputs(
                 inputs_ids,
                 past_key_values,
@@ -606,6 +605,7 @@ def generate(
 
             del idx_next
             progress += 1
+            inputs_ids = inputs_ids_buf.narrow(1, 0, progress)
 
             not_finished = finish.logical_not().to(end_idx.device)
             end_idx.add_(not_finished.int())