project/src/config.rs

@@ -4,7 +4,7 @@ use std::env;
 // Enums
 pub enum OperationMode {
     Training,
-    Inference,
+    Infrence,
 }
 
 // Functions
@@ -15,7 +15,7 @@ pub fn get_operation_mode() -> Option<OperationMode> {
     // Getting operation mode
     match args[1].as_str() {
         "training" => Some(OperationMode::Training),
-        "inference" => Some(OperationMode::Inference),
+        "infrence" => Some(OperationMode::Infrence),
         _ => None,
     }
 }

project/src/main.rs

@@ -13,10 +13,10 @@ fn main() {
 
     // Creating a Neural Network with the Operation Mode
     match operation_mode {
+        None => panic!("Main: `OperationMode` not defined!"),
         Some(mode) => {
             neural = NeuralNetwork::new(mode);
-        },
+        }
-        _ => panic!("Main: `OperationMode` not defined!"),
     }
 
     // Starting the network

project/src/neural.rs

@@ -1,6 +1,6 @@
 // Libraries
 mod data;
-mod inference;
+mod infrence;
 mod model;
 mod training;
 use super::config::OperationMode;
@@ -11,9 +11,6 @@ use burn::{
     optim::AdamConfig,
 };
 
-// Constants
-const MODEL_DIRECTORY: &str = "./out";
-
 // Structures
 pub struct NeuralNetwork {
     mode: OperationMode,
@@ -29,23 +26,25 @@ impl NeuralNetwork {
 
     // Functions
     fn train(&self) {
-        // Creating Backend
         type MyBackend = Wgpu<f32, i32>;
         type MyAutodiffBackend = Autodiff<MyBackend>;
 
         // Create a default Wgpu device
         let device = burn::backend::wgpu::WgpuDevice::default();
 
+        // All the training artifacts will be saved in this directory
+        let artifact_dir = "/tmp/guide";
+
         // Train the model
         training::train::<MyAutodiffBackend>(
-            MODEL_DIRECTORY,
+            artifact_dir,
             training::TrainingConfig::new(model::ModelConfig::new(10, 512), AdamConfig::new()),
             device.clone(),
         );
 
         // Infer the model
-        inference::infer::<MyBackend>(
+        infrence::infer::<MyBackend>(
-            MODEL_DIRECTORY,
+            artifact_dir,
             device,
             burn::data::dataset::vision::MnistDataset::test()
                 .get(42)
@@ -53,13 +52,16 @@ impl NeuralNetwork {
         );
     }
     fn infer(&self) {
-        // Creating Backend
         type MyBackend = Wgpu<f32, i32>;
+
         let device = burn::backend::wgpu::WgpuDevice::default();
 
+        // All the training artifacts are saved in this directory
+        let artifact_dir = "/tmp/guide";
+
         // Infer the model
-        inference::infer::<MyBackend>(
+        infrence::infer::<MyBackend>(
-            MODEL_DIRECTORY,
+            artifact_dir,
             device,
             burn::data::dataset::vision::MnistDataset::test()
                 .get(42)
@@ -71,7 +73,7 @@ impl NeuralNetwork {
         // Switching based on mode
         match self.mode {
             OperationMode::Training => self.train(),
-            OperationMode::Inference => self.infer(),
+            OperationMode::Infrence => self.infer(),
         }
     }
 }