From 5caa633dbc083b2e4fd4e4764976cd5cd2f4e53c Mon Sep 17 00:00:00 2001 From: Guo Yejun Date: Wed, 10 Jan 2024 20:55:57 +0800 Subject: [PATCH] generate.py: do not use args in function main --- generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/generate.py b/generate.py index bb20d6c6..8b291a78 100644 --- a/generate.py +++ b/generate.py @@ -300,7 +300,7 @@ def main( decode_one_token = torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True) # Uncomment to squeeze more perf out of prefill - if args.compile_prefill: + if compile_prefill: prefill = torch.compile(prefill, fullgraph=True, dynamic=True)