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

augmented_images = []

for images, _ in training_set.take(1):
  for i in range(9):
    augmented_image = data_augmentation(np.array(images[i], dtype=jnp.float32))
    augmented_images.append(augmented_image)
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(augmented_images[i].astype("uint8"))
    plt.axis("off")

        __
        __