Update training settings

cnn.py

@@ -52,7 +52,7 @@ class CNNetwork(nn.Module):
             nn.MaxPool2d(kernel_size=2)
         )
         self.flatten = nn.Flatten()
-        self.linear = nn.Linear(128 * 5 * 11, 3)
+        self.linear = nn.Linear(128 * 9 * 31, 3)
         self.softmax = nn.Softmax(dim=1)
 
     def forward(self, input_data):

train.py

@@ -111,9 +111,9 @@ if __name__ == "__main__":
     # instantiating our dataset object and create data loader
     mel_spectrogram = torchaudio.transforms.MelSpectrogram(
         sample_rate=SAMPLE_RATE,
-        n_fft=1024,
+        n_fft=2048,
         hop_length=512,
-        n_mels=64
+        n_mels=128
     )
 
     train_dataset = VoiceDataset(TRAIN_FILE, mel_spectrogram, device)