diff --git a/nanochat/common.py b/nanochat/common.py
index 9462b5b..3ec9992 100644
--- a/nanochat/common.py
+++ b/nanochat/common.py
@@ -128,7 +128,7 @@ def compute_init(device_type="cuda"): # cuda|cpu|mps
         dist.init_process_group(backend="nccl", device_id=device)
         dist.barrier()
     else:
-        device = torch.device(device_type) # cuda|cpu
+        device = torch.device(device_type) # mps|cpu
 
     if ddp_rank == 0:
         logger.info(f"Distributed world size: {ddp_world_size}")