Update DNN architecture and KFold split

app.py

@@ -17,18 +17,22 @@ class DNN(nn.Module):
     def __init__(self, input_size, hidden_size, output_size):
         super(DNN, self).__init__()
         self.fc1 = nn.Linear(input_size, hidden_size)
-        self.relu = nn.ReLU()
-        self.fc2 = nn.Linear(hidden_size, output_size)
+        self.relu1 = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_size, hidden_size)
+        self.relu2 = nn.ReLU()
+        self.fc3 = nn.Linear(hidden_size, output_size)
 
     def forward(self, x):
         x = self.fc1(x)
-        x = self.relu(x)
+        x = self.relu1(x)
         x = self.fc2(x)
+        x = self.relu2(x)
+        x = self.fc3(x)
         return x
 
 
 # Load the model
-model = DNN(input_size=7, hidden_size=128, output_size=1)
+model = DNN(input_size=7, hidden_size=256, output_size=1)
 
 # Initialize the OneHotEncoder
 encoder = OneHotEncoder(handle_unknown="ignore")
@@ -90,7 +94,7 @@ async def train(file: UploadFile = File(...)):
     optimizer = optim.Adam(model.parameters(), lr=0.0001)
 
     # Fit the model
-    kf = KFold(n_splits=4)
+    kf = KFold(n_splits=5)
     accuracies = []
 
     epochs = 50  # Define the number of epochs