text_plain_output_1.png
import matplotlib.pyplot as plt

num_classes = 10
f, ax = plt.subplots(1, num_classes, figsize=(20, 20))
for i in range(0, num_classes):
    sample = x_train[y_train == i][0]
    ax[i].imshow(sample, cmap='gray')
    ax[i].set_title(f'Label: {i}', fontsize=16)