Change model runner

3 files changed, 13 insertions, 9 deletions

diff --git a/Cargo.toml b/Cargo.toml
index c461d72..8de9e1f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,10 +18,12 @@ perf-event = { path = "./perf-event" }
 procfs = { version = "0.17.0", default-features = false }
 dashmap = "6.1.0"
 csv = "1.3.1"
-burn = "0.16.0"
+burn = {version = "0.16.0", features = ["openblas-system","candle"] }
 bincode = "=2.0.0-rc.3"
 bincode_derive = "=2.0.0-rc.3"
-burn-ndarray = "0.16.0"
+#burn-ndarray = {version = "0.16.0", default-features = false, features = ["burn-autodiff", "blas-openblas-system", "std"] }
+#burn-ndarray = {version = "0.16.0", default-features = false, features = ["burn-autodiff"] }
+burn-ndarray = {version = "0.16.0"}
 burn-import = "0.16.0"
 
 [build-dependencies]
diff --git a/power_predictor.py b/power_predictor.py
index cc2a86d..44a7319 100644
--- a/power_predictor.py
+++ b/power_predictor.py
@@ -82,11 +82,11 @@ def analyze_feature_importance(X, y):
 class PowerEstimator(nn.Module):
     def __init__(self, input_size):
         super(PowerEstimator, self).__init__()
-        self.input_lin = nn.Linear(input_size, 16)
+        self.input_lin = nn.Linear(input_size, 4)
         self.relu1 = nn.ReLU()
-        self.lin2 = nn.Linear(16, 4)
+        self.lin2 = nn.Linear(4, 2)
         self.relu2 = nn.ReLU()
-        self.lin3 = nn.Linear(4, 1)
+        self.lin3 = nn.Linear(2, 1)
     
     def forward(self, x):
         x = self.input_lin(x)
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,