plt.figure(figsize=(12, 12))
for img, label in train_dataset.take(1):
    for num in range(16): 
        ax = plt.subplot(4, 4, num + 1)        
        img_to_display = (img[num].numpy() * 255).clip(0, 255).astype("uint8")
        plt.imshow(img_to_display)
        plt.axis("off")  
plt.suptitle("Augmented Dataset")
plt.tight_layout()
plt.show()

        __
        __