From 555356734fd4decdbc3bb655ac91dddee78bf268 Mon Sep 17 00:00:00 2001
From: Artem Bolgar <artem@sesameai.com>
Date: Thu, 29 Feb 2024 18:27:31 +0000
Subject: [PATCH] Minor fix for generate.py

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

diff --git a/generate.py b/generate.py
index 19e1de26..e833e94f 100644
--- a/generate.py
+++ b/generate.py
@@ -310,7 +310,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)