fix local variable 'probs' referenced before assignment issue

teaguesterling · Jul 19, 2023 · 37362e0 · 37362e0
1 parent 4d92db8
commit 37362e0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/llama/generation.py b/llama/generation.py
@@ -167,6 +167,7 @@ def generate(
             # cut to max gen len
             start = 0 if echo else len(prompt_tokens[i])
             toks = toks[start : len(prompt_tokens[i]) + max_gen_len]
+            probs = None
             if logprobs:
                 probs = token_logprobs[i][start : len(prompt_tokens[i]) + max_gen_len]
             # cut to eos tok if any