Just change the builder in this notebook to the one above to use with MNIST images.
Note that ImageClassifier automatically adds a softmax at the end, to get probabilities. If you don’t want that final layer, set finaliser=identity in the ImageClassifier constructor.