import tensorflow_datasets as tfds


def load_nyu():
    builder = tfds.builder('nyu_depth_v2')
    builder.download_and_prepare(download_dir='../nyu')
    return builder.as_dataset(split='train', shuffle_files=True)


if __name__ == '__main__':
    load_nyu()