simplify bool

ahrefs · Apr 22, 2024 · 4e4eba0 · 4e4eba0
1 parent a332ea5
commit 4e4eba0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train_gpt2.py b/train_gpt2.py
@@ -403,7 +403,7 @@ def get_batch():
         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":