revise the model to enhance the training stability

mnist.py

@@ -11,7 +11,7 @@ from lightning.pytorch.callbacks.early_stopping import EarlyStopping
 
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-n", "--n_epochs", type=int, default=200, help="number of epochs of training")
+parser.add_argument("-n", "--n_epochs", type=int, default=1000, help="number of epochs of training")
 parser.add_argument("-b", "--batch", type=int, default=256, help="batch size of training")
 parser.add_argument("-m", "--model", type=str, default='mnist0', help="model to execute")
 opt = parser.parse_args()
@@ -147,14 +147,15 @@ class MNIST_OptAEGV1(MNISTModel):
         self.pool = nn.MaxPool2d(2)
         self.conv0 = nn.Conv2d(1, 2, kernel_size=7, padding=3, bias=False)
         self.lnon0 = OptAEGV1()
-        self.conv1 = nn.Conv2d(2, 2, kernel_size=7, padding=3, bias=False)
+        self.conv1 = nn.Conv2d(2, 2, kernel_size=5, padding=2)
         self.lnon1 = OptAEGV1()
-        self.conv2 = nn.Conv2d(2, 2, kernel_size=7, padding=3, bias=False)
+        self.conv2 = nn.Conv2d(2, 2, kernel_size=5, padding=2)
         self.lnon2 = OptAEGV1()
-        self.conv3 = nn.Conv2d(2, 2, kernel_size=7, padding=3, bias=False)
+        self.conv3 = nn.Conv2d(2, 2, kernel_size=5, padding=2)
         self.lnon3 = OptAEGV1()
-        self.fc = nn.Linear(2 * 3 * 3, 10)
+        self.fc1 = nn.Linear(2 * 3 * 3, 10)
         self.lnon4 = OptAEGV1()
+        self.fc2 = nn.Linear(10, 10, bias=False)
  
     def forward(self, x):
         x = self.conv0(x)
@@ -167,8 +168,9 @@ class MNIST_OptAEGV1(MNISTModel):
         x = self.lnon2(x)
         x = self.pool(x)
         x = th.flatten(x, 1)
-        x = self.fc(x)
+        x = self.fc1(x)
         x = self.lnon4(x)
+        x = self.fc2(x)
         x = F.log_softmax(x, dim=1)
         return x