I got a function implemented as  

def train_model(model, epochs = 1000, batch_size = 64, learning_rate = 0.01):

    # Define loss and optimizer
    criterion = nn.BCELoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
    #scheduler = StepLR(optimizer, step_size=100, gamma=0.90)


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

    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.iloc[i:i + batch_size].values, dtype=torch.float32)

            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())
        
        #scheduler.step() 
        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.values, dtype=torch.float32)
            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)
        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 

Based on this ,I also implemented the following to handle multiple techniques : 

from sklearn.model_selection import KFold
from torch.optim.lr_scheduler import ReduceLROnPlateau

def train_and_evaluate_model(model, learning_rate = 0.01, epochs = 1000, optimization_technique=None, k_splits=None, batch_size=64, patience=None, scheduler_patience=None, **kwargs):
    epoch_train_losses = []
    epoch_train_accuracies = []
    epoch_test_losses = []
    epoch_test_accuracies = []

    if optimization_technique == 'k_fold' and k_splits:
        kfold = KFold(n_splits=k_splits, shuffle=True)
    else:
        kfold = KFold(n_splits=2, shuffle=False)

    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]

        X_train, X_test = torch.tensor(X_train, dtype=torch.float32), torch.tensor(X_test, dtype=torch.float32)
        y_train, y_test = torch.tensor(y_train.values, dtype=torch.float32).view(-1, 1), torch.tensor(y_test.values, dtype=torch.float32).view(-1, 1)

        optimizer = optim.SGD(model.parameters(), lr=learning_rate)
        criterion = nn.BCELoss()
        scheduler = ReduceLROnPlateau(optimizer, 'min', patience=scheduler_patience, verbose=True) if optimization_technique == 'learning_rate_scheduler' else None
        
        best_loss = float('inf')
        stopping_counter = 0
        
        for epoch in range(epochs):
            optimizer.zero_grad()
            y_pred_train = model(X_train)
            loss = criterion(y_pred_train, y_train)
            loss.backward()
            optimizer.step()

            with torch.no_grad():
                y_pred_test = model(X_test)
                test_loss = criterion(y_pred_test, y_test)

            epoch_train_losses.append(loss.item())
            epoch_train_accuracies.extend((y_pred_train > 0.5).float().numpy().flatten().tolist())
            test_accuracy = accuracy_score(y_train, (y_pred_train > 0.5).float()))
            epoch_train_accuracies.append(test_accuracy)
            epoch_test_losses.append(test_loss.item())
            epoch_test_accuracies.extend((y_pred_test > 0.5).float().numpy().flatten().tolist())
            epoch_test_accuracies.append(accuracy_score(y_test, (y_pred_test > 0.5).float()))



            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:
                    break

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

            if optimization_technique == 'k_fold' and k_splits:
                if epoch == 999:
                    break
            if epoch % 100 == 0:
              print(f"Epoch: {epoch+1}/{epochs}, Training Loss: {epoch_train_losses[-1]}, Test Loss: {test_loss.item()}, Training Accuracy: {epoch_train_accuracies[-1]}, Test Accuracy: {test_accuracy}")
    
        if optimization_technique != 'k_fold':
            break

    return epoch_train_losses, epoch_train_accuracies, epoch_test_losses, epoch_test_accuracies


Now train_and_evaluate_model() has 100% of train_model code. So, instead of rewriting the same code again. How can I use the same train_model() code without rewriting in train_and_evaluate_model() like using a lambda function to store the train_model(). Anything like that. Help would be appreciated.