update packages

engine/src/main.rs

@@ -19,12 +19,12 @@ fn select_device() -> Device {
             #[cfg(not(feature = "cuda-backend"))]
             {
                 eprintln!("CUDA requested but not compiled in. Falling back to CPU.");
-                return Device::ndarray();
+                return Device::flex();
             }
         }
         Ok("cpu") => {
             eprintln!("Using CPU device");
-            Device::ndarray()
+            Device::flex()
         }
         Ok("wgpu") => {
             eprintln!("Using WGPU device");
@@ -35,7 +35,7 @@ fn select_device() -> Device {
                 "Unknown CHESS_DRAGON_DEVICE={}. Falling back to CPU.",
                 other
             );
-            Device::ndarray()
+            Device::flex()
         }
     }
 }

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),
-        }
+        })
     }
 }
 
@@ -358,7 +358,7 @@ mod tests {
         #[cfg(feature = "cuda-backend")]
         let device = Device::cuda(DeviceIndex::default());
         #[cfg(not(feature = "cuda-backend"))]
-        let device = Device::ndarray();
+        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