From 5557ac073d86690986b4a73bdd97d8d5388fa4ce Mon Sep 17 00:00:00 2001
From: Lennard Kittner <lennardkittner@icloud.com>
Date: Mon, 31 Mar 2025 17:28:43 +0200
Subject: Add migration message and e core handling in the estimator

---
 src/model.rs | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

(limited to 'src/model.rs')

diff --git a/src/model.rs b/src/model.rs
index ea01d4d..fbb6e0b 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -1,11 +1,10 @@
 use burn::{
-    nn::conv::{Conv2d, Conv2dConfig},
     nn::{Linear, Relu},
     prelude::*,
 };
-use nn::{LeakyReluConfig, LinearConfig};
+use nn::LinearConfig;
 
-use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder, Recorder};
+use burn::record::{FullPrecisionSettings, Recorder};
 use burn_import::pytorch::PyTorchFileRecorder;
 
 //type ArrayBackend = burn_ndarray::NdArray<f32>;
@@ -50,8 +49,20 @@ impl<B: Backend> Net<B> {
     }
 }
 
-/// Load the model from the file in your source code (not in build.rs or script).
-pub fn load_model() -> Net<ArrayBackend> {
+/// Load the p core model from the file in your source code (not in build.rs or script).
+pub fn load_model_p() -> Net<ArrayBackend> {
+    let device = Default::default();
+    let record: NetRecord<ArrayBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
+        .load("./perf.pt".into(), &device)
+        .expect("Failed to decode state");
+
+    Net::<ArrayBackend>::init(&device).load_record(record)
+}
+
+/// Load the e core model from the file in your source code (not in build.rs or script).
+pub fn load_model_e() -> Net<ArrayBackend> {
+    //TODO: load e model
+    println!("Falling back to p model");
     let device = Default::default();
     let record: NetRecord<ArrayBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
         .load("./perf.pt".into(), &device)
-- 
cgit v1.3.1