From b75e860c19aca6929ce94907d4b2a4a03faa69df Mon Sep 17 00:00:00 2001
From: Brett Young <32812705+bdytx5@users.noreply.github.com>
Date: Fri, 14 Mar 2025 05:01:44 -0500
Subject: [PATCH] auto select device based on available hardware

Best to select the device based on what's available (for majority of cases)
---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index de57850510..4c68f9af16 100644
--- a/train.py
+++ b/train.py
@@ -69,7 +69,7 @@
 # DDP settings
 backend = 'nccl' # 'nccl', 'gloo', etc.
 # system
-device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks
+device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'
 dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
 compile = True # use PyTorch 2.0 to compile the model to be faster
 # -----------------------------------------------------------------------------