diff --git a/bird_cnn/Dockerfile b/bird_cnn/Dockerfile
index 4286652..6245a8b 100644
--- a/bird_cnn/Dockerfile
+++ b/bird_cnn/Dockerfile
@@ -1,23 +1,17 @@
 FROM python:3.11-slim
 
-# Set working directory
 WORKDIR /app
 
-# Install system dependencies (optional but common)
 RUN apt-get update && apt-get install -y \
     bash \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy requirements first (better Docker layer caching)
 COPY . .
 
-# Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
 
-RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+RUN pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
 
-# Make sure the startup script is executable
 RUN chmod +x start_server.sh
 
-# Use the script as the container entrypoint
 ENTRYPOINT ["bash", "./start_server.sh"]
\ No newline at end of file
diff --git a/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc b/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc
index 4b07fc7..927c59c 100644
Binary files a/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc and b/bird_cnn/__pycache__/bird_cnn.cpython-314.pyc differ
diff --git a/bird_cnn/bird_cnn.py b/bird_cnn/bird_cnn.py
index a69f897..580e4a6 100644
--- a/bird_cnn/bird_cnn.py
+++ b/bird_cnn/bird_cnn.py
@@ -23,59 +23,49 @@ class  SeparableConvolution(nn.Module):
         x = F.relu(x)
 
         return x
+    
+class SkipBlock(nn.Module):
+    def __init__(self, c_in, c_out, kernel_size=3):
+        super().__init__()
+        self.conv = SeparableConvolution(c_in=c_in, c_out=c_out, kernel_size=kernel_size)
+        self.conv_skip = nn.Sequential(
+            nn.Conv2d(c_in, c_out, 1),
+            nn.BatchNorm2d(c_out)
+        )
+
+    def forward(self, x):
+        return(F.relu(self.conv_skip(x) + self.conv(x)))
 
 
 class Bird_CNN(nn.Module):
     def __init__(self, c_in, c_hidden, c_out):
         super().__init__()
 
-        self.conv_init = nn.Sequential(
+        self.model = nn.Sequential(
             nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
             nn.BatchNorm2d(c_hidden),
-            nn.ReLU()
+            nn.ReLU(),
+
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+            SkipBlock(c_in=c_hidden, c_out=c_hidden),
+
+            SkipBlock(c_in=c_hidden, c_out=c_hidden*2),
+            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
+            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
+            SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
+
+            SeparableConvolution(c_in=c_hidden*2, c_out=c_hidden*4, kernel_size=3),
+
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Flatten(),
+            nn.Linear(c_hidden*4, c_out),
+            nn.Dropout(0.3)
         )
 
-        self.conv_1 = SeparableConvolution(c_in=c_hidden, c_out=c_hidden*2, kernel_size=3)
-        self.conv_skip_1 = nn.Sequential(
-            nn.Conv2d(c_hidden, c_hidden*2, 1),
-            nn.BatchNorm2d(c_hidden*2)
-        )
-
-        self.conv_2 = SeparableConvolution(c_in=c_hidden*2, c_out=c_hidden*4, kernel_size=3)
-        self.conv_skip_2 = nn.Sequential(
-            nn.Conv2d(c_hidden*2, c_hidden*4, 1),
-            nn.BatchNorm2d(c_hidden*4)
-        )
-
-        self.conv_3 = SeparableConvolution(c_in=c_hidden*4, c_out=c_hidden*8, kernel_size=3)
-        self.conv_skip_3 = nn.Sequential(
-            nn.Conv2d(c_hidden*4, c_hidden*8, 1),
-            nn.BatchNorm2d(c_hidden*8)
-        )
-
-        self.conv_4 = SeparableConvolution(c_in=c_hidden*8, c_out=c_hidden*16, kernel_size=3)
-
-        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-        self.flatten = nn.Flatten()
-        self.linear = nn.Linear(c_hidden*16, c_out)
-
-        self.dropout = nn.Dropout(0.3)
-
     def forward(self, x):
-        x = self.conv_init(x)
-
-        x = F.relu(self.conv_skip_1(x) + self.conv_1(x))
-        x = F.relu(self.conv_skip_2(x) + self.conv_2(x))
-        x = F.relu(self.conv_skip_3(x) + self.conv_3(x))
-
-        x = self.conv_4(x)
-
-        x = self.avgpool(x)
-        x = torch.flatten(x, 1)
-
-        x = self.dropout(x)
-
-        return self.linear(x)
+        return self.model(x)
 
 
 def trainCNN(model, optimizer, loss_module, train_data_loader, validation_data_loader, device, num_epochs, SAVE_PATH, save=False):
diff --git a/bird_cnn/main.py b/bird_cnn/main.py
index eeb5c21..24f079e 100644
--- a/bird_cnn/main.py
+++ b/bird_cnn/main.py
@@ -34,13 +34,13 @@ val_size = len(dataset) - train_size
 
 train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
 
-train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
-val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)
+train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)
 
 device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
 print("Using device", device)
 
-model = Bird_CNN(c_in=3, c_hidden=32, c_out=7)
+model = Bird_CNN(c_in=3, c_hidden=64, c_out=7)
 model.to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
 loss_module = nn.CrossEntropyLoss()