device

engine/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2024"
 
 [dependencies]
 chess = "3.2.0"
-burn = { version = "0.22.0-pre.2", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
+burn = { version = "0.22.0-pre.2", features = ["std", "tui", "train", "fusion", "wgpu", "ndarray"], default-features = false }
 rand_distr = "0.6.0"
 rand = "0.10.1"
 ctrlc = "3.5.2"
@@ -16,3 +16,7 @@ tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread", "time
 lru = "0.12"
 futures = "0.3"
 
+[features]
+default = ["cuda-backend"]
+cuda-backend = ["burn/cuda"]
+

engine/src/main.rs

@@ -1,17 +1,44 @@
 #![recursion_limit = "256"]
 
 use burn::optim::AdamConfig;
-use burn::tensor::{Device, DeviceIndex};
+use burn::tensor::Device;
 use engine::mcts::MctsConfig;
 use engine::train::{TrainingConfig, train};
-// fn main() {
-//     type MyBackend = Wgpu<f32, i32>;
-//
-//     let device = Default::default();
-//     let model = ChessModelConfig::new(10, 512).init::<MyBackend>(&device);
-//
-//     println!("{model}");
-// }
+
+#[cfg(feature = "cuda-backend")]
+use burn::tensor::DeviceIndex;
+
+fn select_device() -> Device {
+    match std::env::var("CHESS_DRAGON_DEVICE").as_deref() {
+        Ok("cuda") | Err(_) => {
+            #[cfg(feature = "cuda-backend")]
+            {
+                eprintln!("Using CUDA device");
+                return Device::cuda(DeviceIndex::default());
+            }
+            #[cfg(not(feature = "cuda-backend"))]
+            {
+                eprintln!("CUDA requested but not compiled in. Falling back to CPU.");
+                return Device::ndarray();
+            }
+        }
+        Ok("cpu") => {
+            eprintln!("Using CPU device");
+            Device::ndarray()
+        }
+        Ok("wgpu") => {
+            eprintln!("Using WGPU device");
+            Device::wgpu(Default::default())
+        }
+        Ok(other) => {
+            eprintln!(
+                "Unknown CHESS_DRAGON_DEVICE={}. Falling back to CPU.",
+                other
+            );
+            Device::ndarray()
+        }
+    }
+}
 
 fn main() {
     // Ensure spawned background threads (e.g. CubeCL device servers like DSD-0-0) have sufficient stack space
@@ -21,7 +48,7 @@ fn main() {
         }
     }
 
-    let device = Device::cuda(DeviceIndex::default());
+    let device = select_device();
 
     // Fast-validation config: 400 sims + 48 games ≈ 22 min/iteration on a 3080.
     // Early training with a near-random model benefits more from data diversity

engine/src/net/model.rs

@@ -348,11 +348,17 @@ impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
 mod tests {
     use super::*;
     use burn::optim::{AdamConfig, GradientsParams};
-    use burn::tensor::{Device, DeviceIndex};
+    use burn::tensor::Device;
+
+    #[cfg(feature = "cuda-backend")]
+    use burn::tensor::DeviceIndex;
 
     #[test]
     fn test_train_step() {
+        #[cfg(feature = "cuda-backend")]
         let device = Device::cuda(DeviceIndex::default());
+        #[cfg(not(feature = "cuda-backend"))]
+        let device = Device::ndarray();
         let autodiff_device = device.clone().autodiff();
         let model = ChessModelConfig::init(8, 128, &autodiff_device);
         let mut optim = AdamConfig::new().init();