From 6a477eedbdc8d2c66da84c2fbfcf907ee7e1ba60 Mon Sep 17 00:00:00 2001
From: xiayan0118 <49345397+xiayan0118@users.noreply.github.com>
Date: Mon, 19 Jan 2026 17:19:51 -0800
Subject: [PATCH] fix: pass device_type to compute_init in engine.__main__
 (#451)

When running engine.py directly on non-GPU devices (CPU, MPS),
compute_init() needs the device_type parameter to initialize correctly.
This fixes failures on machines without CUDA support.
---
 nanochat/engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nanochat/engine.py b/nanochat/engine.py
index 7f05eb4..a1ba24c 100644
--- a/nanochat/engine.py
+++ b/nanochat/engine.py
@@ -306,8 +306,8 @@ if __name__ == "__main__":
     """
     import time
     # init compute
-    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init()
     device_type = autodetect_device_type()
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
     autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
 
     # load the model and tokenizer