Hi, I 've implemented two functions train_model() and train_kfold_model(). Have a look at my implementation and fix if you found any mistakes or any other possible errors which may conflict the outputs of all fours optimization techniques.                                                                                                                                                               def train_model(model, X_train, y_train, X_test, y_test, epochs = None, batch_size = None, learning_rate = None, optimization_technique = None,  patience=None, scheduler_patience=None, **kwargs):
    # Define loss and optimizer
    criterion = nn.BCELoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
   
    if optimization_technique == 'learning_rate_scheduler' and scheduler_patience:
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=scheduler_patience, verbose=True)

    # Training segment
    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []

    best_loss = float('inf')
    stopping_counter = 0

    for epoch in range(epochs):
        epoch_train_losses = []
        epoch_y_true = []
        epoch_y_pred = []

        for i in range(0, len(X_train), batch_size):
            X_batch = torch.tensor(X_train[i:i + batch_size], dtype=torch.float)
            y_batch = torch.tensor(y_train[i:i + batch_size].values, dtype=torch.float).view(-1, 1)

            optimizer.zero_grad()
            y_pred = model(X_batch)
            loss = criterion(y_pred, y_batch)
            loss.backward()
            optimizer.step()

            epoch_train_losses.append(loss.item())
            epoch_y_true.extend(y_batch.numpy().flatten().tolist())
            epoch_y_pred.extend((y_pred > 0.5).float().numpy().flatten().tolist())

        train_losses.append(sum(epoch_train_losses) / len(epoch_train_losses))
        train_accuracies.append(accuracy_score(epoch_y_true, epoch_y_pred))

        # Testing segment
        with torch.no_grad():
            X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
            y_test_tensor = torch.tensor(y_test.values, dtype=torch.float32).view(-1, 1)

            test_pred = model(X_test_tensor)
            test_loss = criterion(test_pred, y_test_tensor)
            test_accuracy = accuracy_score(y_test_tensor, (test_pred > 0.5).float())

        test_losses.append(test_loss.item())
        test_accuracies.append(test_accuracy)

        # Early stopping
        if optimization_technique == 'early_stopping' and patience:
            if test_loss < best_loss:
                best_loss = test_loss
                stopping_counter = 0
            else:
                stopping_counter += 1
            if stopping_counter > patience:
                print(f"Early stopping at epoch {epoch+1}/{epochs}")
                break


        # Learning rate scheduler
        if optimization_technique == 'learning_rate_scheduler' and scheduler_patience and scheduler:
            scheduler.step(test_loss)

        if epoch % 100 == 0:
            print(f"Epoch: {epoch+1}/{epochs}, Training Loss: {train_losses[-1]}, Test Loss: {test_loss.item()}, Training Accuracy: {train_accuracies[-1]}, Test Accuracy: {test_accuracy}")

    return train_losses, train_accuracies, test_losses, test_accuracies




def train_kfold_model(model, X_scaled, y, learning_rate = None, epochs = None, k_splits= None, batch_size= None, patience=None, scheduler_patience=None, **kwargs):
    
    results = []
  
    kfold = KFold(n_splits=k_splits, shuffle=True)

    for train_index, test_index in kfold.split(X_scaled):
        X_train, X_test = X_scaled[train_index], X_scaled[test_index]
        y_train, y_test = y[train_index], y[test_index]

        # Call train_model function
        train_losses, train_accuracies, test_losses, test_accuracies = train_model(model, X_train, y_train, X_test, y_test, epochs, batch_size=batch_size, learning_rate=learning_rate, patience=patience, scheduler_patience=scheduler_patience)
        results.append((train_losses, train_accuracies, test_losses, test_accuracies))

    # Process results as needed
    train_losses, train_accuracies, test_losses, test_accuracies = zip(*results)
    train_losses = [item for sublist in train_losses for item in sublist]
    train_accuracies = [item for sublist in train_accuracies for item in sublist]
    test_losses = [item for sublist in test_losses for item in sublist]
    test_accuracies = [item for sublist in test_accuracies for item in sublist]

    return train_losses, train_accuracies, test_losses, test_accuracies

def weights_init(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
        nn.init.zeros_(m.bias)

he_init_model = BetterNNClassifier(input_size, hidden_size, output_size, best_dropout)
he_init_model.apply(weights_init)

# Step 4: Train the base model with four different optimization techniques
train_times = []

start_time = time.time()
train_losses_w, train_accuracies_w, test_losses_w, test_accuracies_w = train_model(he_init_model, X_train, y_train, X_test, y_test, epochs= 1000, batch_size=64, learning_rate=0.01, optimization_technique = 'weight_init')
end_time = time.time()
train_times.append(end_time - start_time)

plot_graphs([train_losses_w, train_accuracies_w, test_losses_w, test_accuracies_w ], technique= 'Weight Initialisation')

start_time = time.time()
train_losses_e, train_accuracies_e, test_losses_e, test_accuracies_e = train_model(best_model,X_train, y_train, X_test, y_test, epochs= 1000, batch_size=64, learning_rate=0.01,  optimization_technique = 'early_stopping', patience=50)
end_time = time.time()
train_times.append(end_time - start_time)

plot_graphs([train_losses_e, train_accuracies_e, test_losses_e, test_accuracies_e], technique = 'Early Stopping')

start_time = time.time()
train_losses_k, train_accuracies_k, test_losses_k, test_accuracies_k = train_kfold_model(best_model, X_scaled, y, epochs= 1000, batch_size=64, learning_rate=0.01, optimization_technique = 'k_fold', k_splits=5)
end_time = time.time()
train_times.append(end_time - start_time)

plot_graphs([train_losses_k, train_accuracies_k, test_losses_k, test_accuracies_k], technique = 'K-Fold')

start_time = time.time()
train_losses_l, train_accuracies_l, test_losses_l, test_accuracies_l = train_model(best_model, X_train, y_train, X_test, y_test, epochs= 1000, batch_size=64, learning_rate=0.01, optimization_technique = 'learning_rate_scheduler', scheduler_patience=50)
end_time = time.time()
train_times.append(end_time - start_time)

plot_graphs([train_losses_l, train_accuracies_l, test_losses_l, test_accuracies_l], technique = 'Learning Rate_Scheduler')


# Step 5: Plot the graphs comparing the test and training accuracy for all optimization methods
optimization_methods = ['Early Stopping', 'k-Fold', 'Learning Rate Scheduler', 'Weight Initialization']

performance_metrics = [max(test_accuracies_e), max(test_accuracies_k), max(test_accuracies_l), max(test_accuracies_w)]

best_optimization_method_index = np.argmax(performance_metrics)
best_optimization_method = optimization_methods[best_optimization_method_index]
best_accuracy = performance_metrics[best_optimization_method_index]

print(f"Best Optimization Method: {best_optimization_method} with test accuracy: {best_accuracy}")


training_time = train_times[best_optimization_method_index]

print(f"Training Time: {training_time:.2f} seconds")