From 957572b08929e88734201933b316b0a7c42b2458 Mon Sep 17 00:00:00 2001
From: Lennard Kittner <lennard@kittner.dev>
Date: Sun, 30 Mar 2025 19:28:25 +0200
Subject: Change model runner

---
 src/model.rs | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

(limited to 'src')

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,
-- 
cgit v1.3.1