Minor fix for generate.py (#117)

pytorch-labs · Mar 11, 2024 · fe84d5b · fe84d5b
1 parent 1f945ab
commit fe84d5b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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)