traing_gpt2.py: Possibility to set device via environment variable

karpathy · Apr 25, 2024 · ad5e912 · ad5e912
1 parent e86d63a
commit ad5e912
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -35,7 +35,7 @@ jobs:
         run: python prepro_tinyshakespeare.py
 
       - name: Train model
-        run: python train_gpt2.py
+        run: DEVICE=cpu python train_gpt2.py
 
       - name: Compile training and testing program
         run: make test_gpt2 train_gpt2

diff --git a/train_gpt2.py b/train_gpt2.py
@@ -339,6 +339,9 @@ def write_tokenizer(enc, filename):
         device = "cuda"
     elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
         device = "mps"
+    # Override with environment variable if set
+    if 'DEVICE' in os.environ:
+        device = os.environ['DEVICE']
     print(f"using device: {device}")
 
     # create a context manager following the desired dtype and device