Explorar el Código

Context length reading configuration file (#540)

刘悦 hace 1 año
padre
commit
79072603b9
Se han modificado 1 ficheros con 3 adiciones y 1 borrados
  1. 3 1
      tools/llama/generate.py

+ 3 - 1
tools/llama/generate.py

@@ -241,7 +241,9 @@ def generate(
 
     codebook_dim = 1 + model.config.num_codebooks
     # create an empty tensor of the expected final shape and fill in the current tokens
-    empty = torch.empty((codebook_dim, max_new_tokens), dtype=dtype, device=device)
+    empty = torch.empty(
+        (codebook_dim, model.config.max_seq_len), dtype=dtype, device=device
+    )
     empty[:, :T] = prompt
     seq = empty
     input_pos = torch.arange(0, T, device=device)