def los_train(df):
    """
    Calculates length of stay (admission to discharge)
    for the training dataset
    Adds len_stay to dataset

    Parameters
    ----------
    df : main dataset
    
    Returns
    -------
    df : main dataset

    """
    df['len_stay'] = (df['DischDateHospProvSpell'] - df['StartDateHospProvSpell']) 
    return df