Updates

engine/src/main.rs

@@ -1,8 +1,7 @@
 #![recursion_limit = "256"]
 
-use burn::backend::Autodiff;
+use burn::backend::{Autodiff, Cuda};
 use burn::optim::AdamConfig;
-use burn_ndarray::NdArray;
 use engine::mcts::MctsConfig;
 use engine::training::train::{train, TrainingConfig};
 // fn main() {
@@ -16,14 +15,14 @@ use engine::training::train::{train, TrainingConfig};
 
 fn main() {
     // type MyBackend = Wgpu<f32, i32>;
-    // type MyBackend = Cuda<f32, i32>;
-    type MyBackend = NdArray<f32, i32>;
+    type MyBackend = Cuda<f32, i32>;
+    // type MyBackend = NdArray<f32, i32>;
     type MyAutodiffBackend = Autodiff<MyBackend>;
     // let device = burn::backend::wgpu::WgpuDevice::default();
-    let device = burn::backend::ndarray::NdArrayDevice::default();
-    // let device = burn::backend::cuda::CudaDevice::default();
+    // let device = burn::backend::ndarray::NdArrayDevice::default();
+    let device = burn::backend::cuda::CudaDevice::default();
 
-    let mcts_config = MctsConfig::new(100, 1.0, 0.05, 0.25);
+    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 
@@ -41,6 +40,7 @@ fn main() {
         mcts_config,
         optimizer: adam_config,
         lr: 2e-4,
+        seed: None,
     };
 
     train::<MyAutodiffBackend>(training_config, device);