from tqdm import tqdm


def get_epochs_with_optional_tqdm(tqdm_mode, nepochs):
    """Get epochs with optional progress bar.

    Args:
        tqdm_mode (str): Progress bar mode.
        nepochs (int): Number of epochs.

    Returns:
        iterable: Epochs.
    """
    if tqdm_mode == "tqdm":

        epochs = tqdm(range(1, nepochs + 1), desc="epoch")
    else:
        epochs = range(1, nepochs + 1)

    return epochs