This gist shows how to fetch batches of data downloaded using tfds.load()

fetching_batches.py

plt.figure(figsize=(12,12)) 

for batch in train.take(1):
    for i in range(9):
        image, label = batch[0][i], batch[1][i]
        plt.subplot(3, 3, i+1)
        plt.imshow(image.numpy())
        plt.title(get_label_name(label.numpy()))
        plt.grid(False)    
# OR
for batch in tfds.as_numpy(train):
    for i in range(9):
        image, label = batch[0][i], batch[1][i]
        plt.subplot(3, 3, i+1)
        plt.imshow(image)
        plt.title(get_label_name(label))
        plt.grid(False)
    # We need to break the loop else the outer loop
    # will loop over all the batches in the training set
    break