Update model.py

model.py

@@ -1,4 +1,3 @@
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -6,13 +5,20 @@ import torch.nn.functional as F
 class AudioCNN(nn.Module):
     def __init__(self, num_classes, input_shape):
         super(AudioCNN, self).__init__()
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        # Convolution layers
+        self.conv1 = nn.Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
         self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+
+        # Dynamically compute the flattened size after conv + pooling
+        with torch.no_grad():
+            dummy_input = torch.zeros(1, *input_shape)  # batch=1
+            x = self.pool(F.relu(self.conv1(dummy_input)))
+            x = self.pool(F.relu(self.conv2(x)))
+            flattened_size = x.numel() // x.size(0)
 
-        pooled_height = input_shape[1] // (2 * 2)
-        pooled_width = input_shape[2] // (2 * 2)
-        self.fc1 = nn.Linear(64 * pooled_height * pooled_width, 128)
+        # Fully connected layers
+        self.fc1 = nn.Linear(flattened_size, 128)
         self.fc2 = nn.Linear(128, num_classes)
 
     def forward(self, x):
@@ -22,3 +28,4 @@ class AudioCNN(nn.Module):
         x = F.relu(self.fc1(x))
         x = self.fc2(x)
         return x
+