title: def show_batch( ) not showing my train images
body: <p>I am preprocessing data in TensorFlow following the instruction on their website: <a href="https://www.tensorflow.org/tutorials/load_data/images" rel="nofollow noreferrer">https://www.tensorflow.org/tutorials/load_data/images</a></p>
  
  <p>However, after turning the images into tensors and assigning each of them to their corresponding label, I cannot plot them.</p>
  
  <p>I load the pair (image, label) in parallel:
  <code>labeled_ds = list_ds.map(process_path, num_parallel_calls=AUTOTUNE)</code>. Then I check for the images shape and corresponding label:</p>
  
  <pre><code>for image, label in labeled_ds.take(1):
    print("Image shape: ", image.numpy().shape)
    print("Label: ", label.numpy())
  </code></pre>
  
  <p>And get
  <code>Image shape:  (80, 80, 3)  Label: [False False True False]</code> as expected.</p>
  
  <p>Then I define the following function to prepare the dataset in batches for training:</p>
  
  <pre><code>def prepare_for_training(ds, cache=True, shuffle_buffer_size=1000):
    ds = ds.batch(100)
    ds = ds.prefetch(buffer_size=AUTOTUNE)
  
    if cache:
      if isinstance(cache, str):
        ds = ds.cache(cache)
      else:
        ds = ds.cache()
  
    ds = ds.shuffle(buffer_size=shuffle_buffer_size)
    ds = ds.repeat()
  
    return ds
  
  
  
  train_ds = prepare_for_training(labeled_ds)
  image_batch, label_batch = next(iter(train_ds))
  </code></pre>
  
  <p>But when I want to display each image with its label using <code>plt.show()</code>, the images don't show. This is how I am doing it:</p>
  
  <pre><code>def show_batch(image_batch, label_batch):
      plt.figure(figsize=(10,10))
      for n in range(25):
          ax = plt.subplot(5,5,n+1)
          plt.imshow(image_batch[n])
          plt.title(CLASS_NAMES[label_batch[n]==1][0].title())
          plt.axis('off')
          return plt.show()
  
  
  show_batch(image_batch.numpy(), label_batch.numpy())
  </code></pre>
  
  <p>Any clues on why my images might not be displaying? </p>