def _batch_size_check(dataset, batch_size):
    """
    Make sure that the `batch_size` size is not bigger than the dataset.
    If it is, it reduces `batch_size` to the lenght of the dataset.
    This is necessary because if PyTorch data_loader receive a batch size bigger than
    the dataset_size it returns an empty generator.

    Parameters
    ----------
    dataset : Dataset
        PyTorch Dataset object.
    batch_size : int
        Size of mini-batches for on the given dataset.

    Returns
    -------
    batch_size : int
        Size of mini-batches for on the given dataset.
        Update to dataset size if needed.
    """

    dataset_size = len(dataset)
    if dataset_size < batch_size:
        print(f"Warning: batch_size({batch_size}) > than dataset_size({dataset_size})."
              "\nSetting batch_size to dataset_size.")
        return dataset_size
    return batch_size