Change model runner

1 files changed, 6 insertions, 4 deletions

diff --git a/src/model.rs b/src/model.rs
index 04eb6c8..ea01d4d 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -8,7 +8,9 @@ use nn::{LeakyReluConfig, LinearConfig};
 use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder, Recorder};
 use burn_import::pytorch::PyTorchFileRecorder;
 
-type ArrayBackend = burn_ndarray::NdArray<f32>;
+//type ArrayBackend = burn_ndarray::NdArray<f32>;
+//pub type ArrayBackend = burn_wgpu::Wgpu;
+pub type ArrayBackend = burn::backend::candle::Candle;
 
 #[derive(Module, Debug)]
 pub struct Net<B: Backend> {
@@ -23,11 +25,11 @@ impl<B: Backend> Net<B> {
     /// Create a new model.
     pub fn init(device: &B::Device) -> Self {
         let input_size = 8;
-        let input_lin = LinearConfig::new(input_size, 16).init(device);
+        let input_lin = LinearConfig::new(input_size, 4).init(device);
         let relu1 = Relu::new();
-        let lin2 = LinearConfig::new(16, 4).init(device);
+        let lin2 = LinearConfig::new(4, 2).init(device);
         let relu2 = Relu::new();
-        let lin3 = LinearConfig::new(4, 1).init(device);
+        let lin3 = LinearConfig::new(2, 1).init(device);
         Self {
             input_lin,
             relu1,