update packages

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 = ["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::flex();
+            }
+        }
+        Ok("cpu") => {
+            eprintln!("Using CPU device");
+            Device::flex()
+        }
+        Ok("wgpu") => {
+            eprintln!("Using WGPU device");
+            Device::wgpu(Default::default())
+        }
+        Ok(other) => {
+            eprintln!(
+                "Unknown CHESS_DRAGON_DEVICE={}. Falling back to CPU.",
+                other
+            );
+            Device::flex()
+        }
+    }
+}
 
 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

@@ -3,11 +3,11 @@ use burn::data::dataloader::batcher::Batcher;
 use burn::nn::conv::Conv2dConfig;
 use burn::nn::loss::{MseLoss, Reduction};
 use burn::nn::{BatchNorm, BatchNormConfig, LinearConfig, PaddingConfig2d};
-use burn::tensor::activation::log_softmax;
 use burn::tensor::Transaction;
+use burn::tensor::activation::log_softmax;
 use burn::train::{InferenceStep, ItemLazy, TrainOutput, TrainStep};
 use burn::{
-    nn::{conv::Conv2d, Linear, Relu},
+    nn::{Linear, Relu, conv::Conv2d},
     prelude::*,
 };
 /*
@@ -209,7 +209,7 @@ pub struct ChessOutput {
 }
 
 impl ItemLazy for ChessOutput {
-    fn sync(self) -> Self {
+    fn sync(self) -> Result<Self, burn::tensor::ExecutionError> {
         let [
             policy_logits,
             value,
@@ -232,7 +232,7 @@ impl ItemLazy for ChessOutput {
 
         let device = &Default::default();
 
-        ChessOutput {
+        Ok(ChessOutput {
             policy_logits: Tensor::from_data(policy_logits, device),
             value: Tensor::from_data(value, device),
             policy_targets: Tensor::from_data(policy_targets, device),
@@ -240,7 +240,7 @@ impl ItemLazy for ChessOutput {
             policy_loss: Tensor::from_data(policy_loss, device),
             value_loss: Tensor::from_data(value_loss, device),
             loss: Tensor::from_data(loss, device),
-        }
+        })
     }
 }
 
@@ -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::flex();
         let autodiff_device = device.clone().autodiff();
         let model = ChessModelConfig::init(8, 128, &autodiff_device);
         let mut optim = AdamConfig::new().init();

engine/src/train.rs

@@ -5,7 +5,6 @@ use crate::net::model::{
 };
 use crate::worker::{NNRequest, NNServer, NNServerCmd, Worker, WorkerCmd, WorkerData};
 use burn::data::dataloader::batcher::Batcher;
-use burn::module::AutodiffModule;
 use burn::optim::{AdamConfig, GradientsParams};
 use burn::prelude::{Device, Module};
 use burn::store::ModuleRecord;
@@ -272,18 +271,18 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
                 model = optim.step(training_config.lr, model, grads);
 
                 if step == training_config.num_train_steps - 1 {
-                    loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+                    loss_val = output.loss.clone().into_data().try_to_vec_as::<f32>().unwrap()[0];
                     p_loss_val = output
                         .policy_loss
                         .clone()
                         .into_data()
-                        .to_vec::<f32>()
+                        .try_to_vec_as::<f32>()
                         .unwrap()[0];
                     v_loss_val = output
                         .value_loss
                         .clone()
                         .into_data()
-                        .to_vec::<f32>()
+                        .try_to_vec_as::<f32>()
                         .unwrap()[0];
                 }
             }

engine/src/worker.rs

@@ -411,8 +411,8 @@ impl NNServer {
             let (p, v) = self.model.forward(batched_tensor);
             batches_ran += 1;
 
-            let p_vec = p.into_data().to_vec::<f32>().unwrap();
-            let v_vec = v.into_data().to_vec::<f32>().unwrap();
+            let p_vec = p.into_data().try_to_vec_as::<f32>().unwrap();
+            let v_vec = v.into_data().try_to_vec_as::<f32>().unwrap();
 
             const POLICY_DIM: usize = 73 * 8 * 8; // 4672