'Invalid shape (28, 28, 1) for image data'

Question

[enter image description here][1]

(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()

train_images = trai