Fix warnings and general cleanup

1 files changed, 6 insertions, 8 deletions

diff --git a/src/model.rs b/src/model.rs
index 6dbde1f..eb62c35 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -7,9 +7,7 @@ use nn::LinearConfig;
 use burn::record::{FullPrecisionSettings, Recorder};
 use burn_import::pytorch::PyTorchFileRecorder;
 
-//type ArrayBackend = burn_ndarray::NdArray<f32>;
-//pub type ArrayBackend = burn_wgpu::Wgpu;
-pub type ArrayBackend = burn::backend::candle::Candle;
+pub type BurnBackend = burn::backend::candle::Candle;
 
 #[derive(Module, Debug)]
 pub struct Net<B: Backend> {
@@ -44,17 +42,17 @@ impl<B: Backend> Net<B> {
         let x = self.relu1.forward(x);
         let x = self.lin2.forward(x);
         let x = self.relu2.forward(x);
-        let x = self.lin3.forward(x);
-        x
+
+        self.lin3.forward(x)
     }
 }
 
 /// Load the p core model from the file in your source code (not in build.rs or script).
-pub fn load_model(path: &str) -> Net<ArrayBackend> {
+pub fn load_model(path: &str) -> Net<BurnBackend> {
     let device = Default::default();
-    let record: NetRecord<ArrayBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
+    let record: NetRecord<BurnBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
         .load(path.into(), &device)
         .expect("Failed to decode state");
 
-    Net::<ArrayBackend>::init(&device).load_record(record)
+    Net::<BurnBackend>::init(&device).load_record(record)
 }