added loading and testing the model again

architectures/Model.py

@@ -62,9 +62,7 @@ class Model(ABC):
 
 
     def save(self, filename=None):
-        """
-        Saves the model state_dict. Creates the directory if it doesn't exist.
-        """
+        
         save_dir = Path("trained_models")
         save_dir.mkdir(parents=True, exist_ok=True)
 
@@ -72,8 +70,29 @@ class Model(ABC):
         if filename is None:
             filename = f"{self.__class__.__name__.lower()}.pth"
         
+        if not filename.endswith('.pth'):
+            filename += '.pth'
+
         save_path = save_dir / filename
         torch.save(self.model.state_dict(), save_path)
+        print(f'Model saved to {save_path}')
+
+
+
+    def load(self, arch):
+ 
+        file_path = Path("trained_models") / f'{arch.name.lower()}.pth'
+        # does file exist
+        if not file_path.exists():
+            raise FileNotFoundError(f'No checkpoint found at: {file_path}')
+            
+        # Load the weights
+        state_dict = torch.load(file_path, map_location=self.device, weights_only=True)
+        
+        self.model.load_state_dict(state_dict)
+        self.model.to(self.device)
+        print(f'Model loaded from {file_path}')
+
 
 
     # Using the factory patern here