From f0e14b634124590600ce380875c3f7ec290dd69a Mon Sep 17 00:00:00 2001
From: Daniel Strobusch <1847260+dastrobu@users.noreply.github.com>
Date: Mon, 18 Dec 2023 19:11:51 +0100
Subject: [PATCH] fix renamed arg (#140)

---
 llama/llama.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llama/llama.py b/llama/llama.py
index af98685d..ad6fd8ce 100644
--- a/llama/llama.py
+++ b/llama/llama.py
@@ -294,7 +294,7 @@ def few_shot_generate(args):
                 mx.eval(token)
                 prompt_processing = toc("Prompt processing", start)
 
-            if len(tokens) >= args.num_tokens:
+            if len(tokens) >= args.max_tokens:
                 break
 
             mx.eval(tokens)