def get_predict_dataset(distribution):
    inputs = np.zeros((10, 3), dtype=np.float32)
    dataset = dataset_ops.Dataset.from_tensor_slices(inputs)
    dataset = dataset.repeat(100)
    dataset = batch_wrapper(dataset, 10, distribution)
    return dataset
def get_predict_dataset(distribution):
    inputs = np.zeros((10, 3), dtype=np.float32)
    dataset = dataset_ops.Dataset.from_tensor_slices(inputs)
    dataset = dataset.repeat(100)
    dataset = batch_wrapper(dataset, 10, distribution)
    return dataset