Fashion MNIST in TensorFlow

#!/usr/bin/python3

import tensorflow as tf

class MnCb(tf.keras.callbacks.Callback):
        def on_epoch_end(self, epoch, logs={}):
                if (logs.get('loss') < 0.4):
                        print('\nLoss is LESS\n')
                        self.model.stop_training = True

cbs = MnCb()

num_cats = 10

fmnist = tf.keras.datasets.fashion_mnist

(xtrain, ytrain), (xtest, ytest) = fmnist.load_data()

xtrain = xtrain / 255.
xtest = xtest / 255.

#ytrain = tf.keras.utils.to_categorical(ytrain, num_cats)
#ytest = tf.keras.utils.to_categorical(ytest, num_cats)

model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=[28,28,1]),
        tf.keras.layers.Conv2D(filters=32, kernel_size=[3,3], activation='relu'),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Conv2D(filters=64, kernel_size=[3,3], activation='relu'),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(units=num_cats, activation='softmax')
])

model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy'])

model.fit(xtrain, ytrain, epochs=5, batch_size=128, callbacks=[cbs])

model.evaluate(xtest, ytest)