Update : Modifier details in the model

main.py

@@ -78,7 +78,7 @@
     print(model_scratch)
 
     model_scratch = train_class.train(10, loaders_scratch, model_scratch, optimizer_scratch, 
-                      criterion_scratch, use_cuda, './model/model_scratch.pt') # You can rename this file to save different check point
+                      criterion_scratch, use_cuda, './model/model_cnn_2.pt') # You can rename this file to save different check point
 
     # load the model that got the best validation accuracy
     # model_scratch.load_state_dict(torch.load('./model/model_scratch.pt')) # uncomment only to load the saved model

model.py

@@ -10,8 +10,8 @@ def __init__(self, n_feature, output_size):
 
         ## Define layers of a CNN
         self.conv1 = nn.Conv2d(3, n_feature*2**1, 3, padding=1, stride=2)
-        self.conv2 = nn.Conv2d(32, n_feature*2**2, 3, padding=1, stride=2)
-        self.conv3 = nn.Conv2d(64, n_feature*2**3, 3, padding=1)
+        self.conv2 = nn.Conv2d(n_feature*2**1, n_feature*2**2, 3, padding=1, stride=2)
+        self.conv3 = nn.Conv2d(n_feature*2**2, n_feature*2**3, 3, padding=1)
 
         # max pooling layer
         self.pool = nn.MaxPool2d(2, 2)
@@ -45,9 +45,9 @@ def __init__(self, n_feature, output_size):
 
         ## Define layers of a CNN
         self.conv1 = nn.Conv2d(3, n_feature*2**1, 3, padding=1, stride=2)
-        self.conv2 = nn.Conv2d(32, n_feature*2**2, 3, padding=1, stride=2)
-        self.conv3 = nn.Conv2d(64, n_feature*2**3, 3, padding=1)
-        self.conv4 = nn.Conv2d(128, n_feature*2**4, 3, padding=1, stride=2)
+        self.conv2 = nn.Conv2d(n_feature*2**1, n_feature*2**2, 3, padding=1, stride=2)
+        self.conv3 = nn.Conv2d(n_feature*2**2, n_feature*2**3, 3, padding=1)
+        self.conv4 = nn.Conv2d(n_feature*2**3, n_feature*2**4, 3, padding=1, stride=2)
 
         # max pooling layer
         self.pool = nn.MaxPool2d(2, 2)