Add migration message and e core handling in the estimator

1 files changed, 16 insertions, 5 deletions

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)