for i, (img, label) in enumerate(train_data.take(16)):
    ax = plt.subplot(4,4,i+1)
    plt.imshow(img)
    plt.title(label.numpy())

        __
        __