diff --git a/.gitignore b/.gitignore
index f514b74..037bf3a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,14 @@
-# Created by venv; see https://docs.python.org/3/library/venv.html
-*
+# Virtual Environment (the folders Git saw)
+bin/
+lib/
+share/
+pyvenv.cfg
+include/
+
+# Data and Models
+data/
+trained_models/
+
+# Python cache
+__pycache__/
+*.py[cod]
diff --git a/Tune.py b/Tune.py
index 2f2a164..402e658 100644
--- a/Tune.py
+++ b/Tune.py
@@ -8,7 +8,7 @@ from IdentitySubset import IdentitySubset
 from architectures.Model import Model, Architecture
 
 # numbre of classes 
-CLASS_SIZE = 20
+CLASS_SIZE = 50
 # batch 
 BATCH_SIZE = 16
 
@@ -34,7 +34,9 @@ RESOLUTION = 224
 #  - DENSENET121
 #  - INCEPTION
 #  - GOOGLENET
-arch = Architecture.EFFICIENTNET
+#  - EFFICIENTNET
+#  - SHUFFLENET
+arch = Architecture.GOOGLENET
 
 # DATA PREPARATION
 # load data set and prepare
@@ -93,9 +95,12 @@ model.train(
     rate = LR_RATE)
 
 # save.
+model.save(filename=arch.name.lower())
+'''
 torch.save(
     model.get().state_dict(), 
-    f'{arch.name.lower()}.pth')
+    f'trained/{arch.name.lower()}.pth'
+)'''
 
 # done tuning
 print('Model saved!')
diff --git a/architectures/DenseNet121.py b/architectures/DenseNet121.py
index 71a2576..68a6735 100644
--- a/architectures/DenseNet121.py
+++ b/architectures/DenseNet121.py
@@ -12,4 +12,4 @@ class DenseNet121(Model):
         num_ftrs = m.classifier.in_features
         m.classifier = nn.Linear(num_ftrs, self.size)
         
-        return m.to(self.device)
\ No newline at end of file
+        return m
\ No newline at end of file
diff --git a/architectures/EfficentNet.py b/architectures/EfficentNet.py
new file mode 100644
index 0000000..9fbfc97
--- /dev/null
+++ b/architectures/EfficentNet.py
@@ -0,0 +1,19 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class EfficientNet(Model):
+    
+    def get(self):
+
+        m = models.efficientnet_b1(weights=models.EfficientNet_B1_Weights.DEFAULT)
+
+        # Unfreeze the last block for a lighter touch
+        for param in m.features[-1].parameters(): param.requires_grad = True
+
+        # Standard classifier fix
+        m.classifier[1] = nn.Linear(m.classifier[1].in_features, self.size)
+        return m
\ No newline at end of file
diff --git a/architectures/GoogleNet.py b/architectures/GoogleNet.py
new file mode 100644
index 0000000..33a535e
--- /dev/null
+++ b/architectures/GoogleNet.py
@@ -0,0 +1,14 @@
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class GoogleNet(Model):
+    
+    def get(self):
+
+        m = models.googlenet(weights =  models.GoogLeNet_Weights.DEFAULT)
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m
\ No newline at end of file
diff --git a/architectures/Inception.py b/architectures/Inception.py
index aff33b8..b55997f 100644
--- a/architectures/Inception.py
+++ b/architectures/Inception.py
@@ -10,12 +10,12 @@ from architectures.Model import Model
 
 class Inception(Model):
     def get(self):
-        model = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT)
+        m = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT)
         #for param in model.parameters():
         #    param.requires_grad = False
-        model.fc = nn.Linear(model.fc.in_features, self.size)
-        model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, self.size)
-        return model.to(self.device)
+        m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, self.size)
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m
 
     def train(self, epochs, loader, rate):
         # Override because Inception returns a tuple (main, aux)
diff --git a/architectures/Model.py b/architectures/Model.py
index 53e7aa5..1c20997 100644
--- a/architectures/Model.py
+++ b/architectures/Model.py
@@ -6,17 +6,18 @@ import torch.optim as optim
 import time
 import numpy as np
 from sklearn.metrics import classification_report
+from pathlib import Path
 
 class Model(ABC):
     def __init__(self, device, size):
         self.device = device
         self.size = size
-        self.model = self.get()
+        self.model = self.get().to(self.device)
 
     @abstractmethod
     def get(self):
         # return the model
-        return self.model
+        pass
 
     def train(self, epochs, loader, rate):
         criterion = nn.CrossEntropyLoss()
@@ -60,6 +61,21 @@ class Model(ABC):
         print(classification_report(all_labels, all_preds, zero_division=0))
 
 
+    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)
+
+        # 2. Determine filename (Default to class name if not provided)
+        if filename is None:
+            filename = f"{self.__class__.__name__.lower()}.pth"
+        
+        save_path = save_dir / filename
+        torch.save(self.model.state_dict(), save_path)
+
+
     # Using the factory patern here
     @staticmethod
     def create(arch, device, size):
@@ -74,8 +90,8 @@ class Model(ABC):
             
             # ResNet50
             case Architecture.RESNET50:
-                from architectures.ResNet18 import ResNet18
-                return ResNet18(device, size)
+                from architectures.ResNet50 import ResNet50
+                return ResNet50(device, size)
             
             # INCEPTION
             case Architecture.INCEPTION:
@@ -94,6 +110,14 @@ class Model(ABC):
             case Architecture.EFFICIENTNET:
                 from architectures.EfficentNet import EfficientNet
                 return EfficientNet(device, size)
+            #ShuffleNet
+            case Architecture.SHUFFLENET:
+                from architectures.ShuffleNet import ShuffleNet
+                return ShuffleNet(device, size)
+            # wide ResNet
+            case Architecture.WIDE_RESNET:
+                from architectures.WideResNet import WideResNet
+                return WideResNet(device, size)
             case _:
                 raise ValueError(f"Unknown model: {arch}")
 
@@ -106,4 +130,6 @@ class Architecture(Enum):
     INCEPTION = auto()
     DENSENET121 = auto()
     GOOGLENET = auto()
-    EFFICIENTNET = auto()
\ No newline at end of file
+    EFFICIENTNET = auto()
+    SHUFFLENET = auto()
+    WIDE_RESNET = auto()
\ No newline at end of file
diff --git a/architectures/ResNet18.py b/architectures/ResNet18.py
index e0c1d9d..3c11924 100644
--- a/architectures/ResNet18.py
+++ b/architectures/ResNet18.py
@@ -11,12 +11,12 @@ class ResNet18(Model):
         m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
 
         # freez all layers
-        #for param in m.parameters():
-        #    param.requires_grad = False
+        for param in m.parameters():
+            param.requires_grad = False
 
         # unfreez the last two
-        #for param in m.layer3.parameters(): param.requires_grad = True
-        #for param in m.layer4.parameters(): param.requires_grad = True
+        for param in m.layer3.parameters(): param.requires_grad = True
+        for param in m.layer4.parameters(): param.requires_grad = True
         
         m.fc = nn.Linear(m.fc.in_features, self.size)
-        return m.to(self.device)
\ No newline at end of file
+        return m
\ No newline at end of file
diff --git a/architectures/ResNet50.py b/architectures/ResNet50.py
index a039ef1..74a4df4 100644
--- a/architectures/ResNet50.py
+++ b/architectures/ResNet50.py
@@ -6,6 +6,18 @@ from torchvision import models
 from architectures.Model import Model
 
 class ResNet50(Model):
+    # NOTE:
+    # This model had it's best performance with the following configs
+    # numbre of classes 
+    # CLASS_SIZE = 20
+    # BATCH_SIZE = 16
+    # SAMPLE_SIZE = 30
+    # TRAINING_SMPLE = 28
+    # LR_RATE = 0.0001
+    # EPOCHS = 15
+    # RESOLUTION = 224
+    # NOTE: But it may be a one time thing. 
+    # because testing again didn't repeat 
     
     def get(self):
         m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
@@ -15,8 +27,10 @@ class ResNet50(Model):
             param.requires_grad = False
 
         # unfreez the last two
+        # NOTE: Freezing everything and unfrizing the last 3 yeilded the best performance
+        for param in m.layer2.parameters(): param.requires_grad = True
         for param in m.layer3.parameters(): param.requires_grad = True
         for param in m.layer4.parameters(): param.requires_grad = True
 
         m.fc = nn.Linear(m.fc.in_features, self.size)
-        return m.to(self.device)
\ No newline at end of file
+        return m
\ No newline at end of file
diff --git a/architectures/ShuffleNet.py b/architectures/ShuffleNet.py
new file mode 100644
index 0000000..fc2a52d
--- /dev/null
+++ b/architectures/ShuffleNet.py
@@ -0,0 +1,17 @@
+
+
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class ShuffleNet(Model):
+    
+    def get(self):
+        m = models.shufflenet_v2_x1_0(weights=models.ShuffleNet_V2_X1_0_Weights.DEFAULT)
+        
+        num_ftrs = m.fc.in_features
+        m.fc = nn.Linear(num_ftrs, self.size)
+        return m
\ No newline at end of file
diff --git a/architectures/WideResNet.py b/architectures/WideResNet.py
new file mode 100644
index 0000000..b645096
--- /dev/null
+++ b/architectures/WideResNet.py
@@ -0,0 +1,15 @@
+
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class WideResNet(Model):
+    
+    def get(self):
+        # wide_resnet50_2 is a common high-performance choice
+        m = models.wide_resnet50_2(weights=models.Wide_ResNet50_2_Weights.DEFAULT)
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m
\ No newline at end of file