def train_model(model, X_train, y_train, epochsPerTrial, nRounds, batchSize,
                saveResults=True, 
                saveFileName=''):
    """
    Function to train the NBGNet model.
    Args:
        model: model to be trained
        X_train: Dictionary, where single-trial input data is stored in corresponding trial as key
        y_train: Dictionary, where single-trial output data is stored in corresponding trial as key
        epochsPerTrial: Integer, number of epochs per trial for model training    
        nRounds: Integer, number of rounds for model training     
        batchSize: Integer, batch size for model training  
        saveResults: Binary, whether save the trained weight or not    
        saveFileName: String, path of the saved .h5 file      
        
    Returns:
        history: training history         
    """
    
    history = {}
    for r in range(nRounds):
        for tr in range(len(X_train.keys())):
            history[(r, tr)] = model.fit(X_train[tr], y_train[tr], batch_size=batchSize, 
                                         epochs=epochsPerTrial)
    
    if saveResults:
        model.save_weights(saveFileName)
    
    return history