Merge pull request meta-llama#403 from tianweidut/bug/fix-probs-refer…

…ence bug: fix local variable 'probs' referenced before assignment typo
vkrit · Jul 19, 2023 · cc8a000 · cc8a000
2 parents 4d92db8 + 37362e0
commit cc8a000
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