Reuse cache for different generation size

fishaudio · Dec 20, 2023 · 0dafcf3 · 0dafcf3
1 parent 95d90c8
commit 0dafcf3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tools/llama/generate.py b/tools/llama/generate.py
@@ -163,7 +163,7 @@ def decode_n_tokens(
     **sampling_kwargs,
 ):
     previous_tokens = torch.zeros(
-        (model.config.num_codebooks + 1, num_new_tokens),
+        (model.config.num_codebooks + 1, model.config.max_seq_len),
         dtype=torch.int,
         device=cur_token.device,
     )