fixes

engine/Cargo.toml

@@ -6,7 +6,6 @@ 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-ndarray = "0.21.0"
 rand_distr = "0.6.0"
 rand = "0.10.1"
 ctrlc = "3.5.2"

engine/src/main.rs

@@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig};
 fn main() {
     let device = Device::cuda(DeviceIndex::default());
 
-    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
+    let mcts_config = MctsConfig::new(100, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 
@@ -37,7 +37,7 @@ fn main() {
         seed: None,
     };
 
-    train(training_config, &device);
+    train(training_config, device);
 
     // let mut board_state: Option<BoardState> = Some(BoardState::new(
     //     Board::default(),

engine/src/train.rs

@@ -6,7 +6,6 @@ use crate::net::model::{
 use burn::data::dataloader::batcher::Batcher;
 use burn::module::{AutodiffModule, Module};
 use burn::optim::{AdamConfig, GradientsParams, Optimizer};
-// use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
 use burn::prelude::Device;
 use burn::store::ModuleRecord;
 use chess::{ChessMove, Color};
@@ -38,20 +37,32 @@ pub struct TrainingConfig {
     pub seed: Option<u64>,
 }
 
-pub fn train(training_config: TrainingConfig, device: &Device) {
+pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
+    let device = device.into();
+    if training_config.seed.is_some() {
+        device.seed(training_config.seed.unwrap());
+    }
+    let autodiff_device = device.clone().autodiff();
+
     let model_path = format!("artifacts/{}", training_config.model_name.as_str());
     let metadata_path = format!("{}.json", model_path);
     println!("Creating model...");
     let mut model: ChessModel = ChessModelConfig::init(
         training_config.num_blocks,
         training_config.hidden_channels,
-        device,
+        &autodiff_device,
     );
+    // Initialize optimizer once so state (moments) persist across steps
+    let mut optim = training_config.optimizer.init();
+
     if training_config.load_model {
         println!("Loading model {}...", model_path);
         // Load model in full precision from MessagePack file
         let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
-        model = model.load_record(record)
+        model = model.load_record(record);
+        optim = optim
+            .load(format!("{model_path}_optim"))
+            .expect("Failed to load optimizer");
     }
 
     let train = Arc::new(AtomicBool::new(true));
@@ -92,9 +103,6 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
     println!("using seed: {}", seed);
     let mut rng = SmallRng::seed_from_u64(seed);
 
-    // Initialize optimizer once so state (moments) persist across steps
-    let mut optim = training_config.optimizer.init();
-
     println!("Starting training...");
     while train.load(Ordering::Relaxed) {
         println!("Iteration: {}", iter);
@@ -116,11 +124,7 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
                 // println!("{}", before.elapsed().as_millis());
                 episode_buffer.push(results);
 
-                let temp = if ply_count < 30 {
-                    1.0
-                } else {
-                    0.0
-                };
+                let temp = if ply_count < 30 { 1.0 } else { 0.0 };
 
                 let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
 
@@ -201,8 +205,18 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
         let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
 
         let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
-        let p_loss_val = output.policy_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
-        let v_loss_val = output.value_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+        let p_loss_val = output
+            .policy_loss
+            .clone()
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
+        let v_loss_val = output
+            .value_loss
+            .clone()
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
 
         println!(
             "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
@@ -254,6 +268,10 @@ pub fn train(training_config: TrainingConfig, device: &Device) {
         .into_record()
         .save(&model_path)
         .expect("Should be able to save the model");
+
+    optim
+        .save(format!("{model_path}_optim"))
+        .expect("Should be able to save the optim");
     println!("Model saved in {:?}, exiting training.", model_path);
     return;
 }