Implement energy model for e cores

1 files changed, 2 insertions, 14 deletions

diff --git a/src/model.rs b/src/model.rs
index fbb6e0b..6dbde1f 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -50,22 +50,10 @@ impl<B: Backend> Net<B> {
 }
 
 /// 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> {
+pub fn load_model(path: &str) -> 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)
+        .load(path.into(), &device)
         .expect("Failed to decode state");
 
     Net::<ArrayBackend>::init(&device).load_record(record)