diff --git a/cs2109s/labs/final/final.py b/cs2109s/labs/final/final.py
index 23a8512..8b11857 100644
--- a/cs2109s/labs/final/final.py
+++ b/cs2109s/labs/final/final.py
@@ -21,31 +21,30 @@ import os
 class CNN3D(nn.Module):
     def __init__(self):
         super(CNN3D, self).__init__()
-        self.conv1 = nn.Conv3d(1, 16, 2, 1, 2)
-        self.batchnorm3d = nn.BatchNorm3d(16)
-
-        self.dropout = nn.Dropout(0.5)
-        self.mp3d = nn.AvgPool3d(2)
-        self.relu = nn.LeakyReLU()
-
-        self.lstm = nn.LSTM(5184, 64, 1, batch_first=True)
-        self.fc2 = nn.Linear(64, 6)
+        self.conv1 = nn.Conv3d(1, 32, kernel_size=3, stride=1, padding=1)
+        self.conv2 = nn.Conv3d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.batchnorm = nn.BatchNorm3d(32)
+        self.relu = nn.ReLU()
+        self.maxpool = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.fc1 = nn.Linear(1024, 256)  # Calculate input size based on output from conv3
+        self.fc2 = nn.Linear(256, 6)
 
     def forward(self, x):
         x = self.conv1(x)
         x = self.relu(x)
-        x = self.batchnorm3d(x)
-        x = self.mp3d(x)
-        x = self.dropout(x)
+        x = self.maxpool(x)
+        x = self.batchnorm(x)
+        x = self.conv2(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
 
-        x = x.view(-1, 5184)
-
-        x, _ = self.lstm(x)
+        x = x.view(x.size(0), -1)  # Flatten features for fully connected layers
+        x = self.fc1(x)
         x = self.relu(x)
         x = self.fc2(x)
-        return torch.softmax(x, dim=1)
+        return x
 
-def train(model, criterion, optimizer, loader, epochs=10):
+def train(model, criterion, optimizer, loader, epochs=5):
     for epoch in range(epochs):
         for idx, (inputs, labels) in enumerate(loader):
             optimizer.zero_grad()
@@ -74,7 +73,7 @@ def process_data(X, y):
     # Get the indices of each class
     indices = [np.argwhere(y == i).squeeze(1) for i in range(6)]
     # Get the number of samples to take for each class
-    num_samples_to_take = 1500
+    num_samples_to_take = 600
     # Get the indices of the samples to take
     indices_to_take = [np.random.choice(indices[i], num_samples_to_take, replace=True) for i in range(6)]
     # Concatenate the indices
@@ -94,8 +93,8 @@ class Model():
     def fit(self, X, y):
         X, y = process_data(X, y)
         train_dataset = torch.utils.data.TensorDataset(X, y)
-        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
-        train(self.model, self.criterion, self.optimizer, train_loader, 5)
+        train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)
+        train(self.model, self.criterion, self.optimizer, train_loader, 10)
 
     def predict(self, X):
         self.model.eval()