def split_dataset(dataset, train_prop=0.8, val_prop=0.2):
    """
    This function takes in the loaded TFRecordDataset, and builds training, validation
    and test TFRecordDataset objects. The test_prop is automatically set up to be equal to
    1 - (train_prop + val_prop).
    """

    dataset_size = sum(1 for _ in dataset)
    train_size = int(train_prop * dataset_size)
    val_size = int(val_prop * dataset_size)
    
    train_dataset = dataset.take(train_size)
    remaining_dataset = dataset.skip(train_size)
    val_dataset = remaining_dataset.take(val_size)
    test_dataset = remaining_dataset.skip(val_size)
    return train_dataset, val_dataset, test_dataset