Changes are made

model.py

@@ -34,7 +34,6 @@ class ConvModel(nn.Module):
         
         self.classifier = nn.Sequential(
                                 nn.Flatten(),
-                                nn.Dropout(0.5),  # Dropout layer to reduce overfitting
                                 nn.Linear(in_features = 40960, out_features = output_shape))