From 4e4eba04410a50894846f6a64be1be2847313d0a Mon Sep 17 00:00:00 2001
From: Andrej Karpathy <andrej.karpathy@gmail.com>
Date: Mon, 22 Apr 2024 00:19:00 +0000
Subject: [PATCH] simplify bool

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

diff --git a/train_gpt2.py b/train_gpt2.py
index 0e93b8c..df65aec 100644
--- a/train_gpt2.py
+++ b/train_gpt2.py
@@ -403,7 +403,7 @@ if __name__ == "__main__":
         write_model(model, "gpt2_124M.bin")
         write_state(model, x, y, logits, loss, "gpt2_124M_debug_state.bin")
 
-    use_fused = True if device == "cuda" else False
+    use_fused = device == "cuda" # only works on CUDA (?)
     optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, fused=use_fused)
     timings = []
     if device == "cuda":