added more models for testing

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!')