MNIST Code in TensorFlow

#!/usr/bin/python3

import tensorflow as tf

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

cbs = MCb()

mnist = tf.keras.datasets.mnist

(xtrain, ytrain), (xtest, ytest) = mnist.load_data()
xtrain = xtrain / 255.
xtest = xtest / 255.
ytrain = tf.keras.utils.to_categorical(ytrain, 10)
ytest = tf.keras.utils.to_categorical(ytest, 10)

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

model.compile(
        optimizer='sgd',
        loss='categorical_crossentropy',
        metrics=['accuracy'])

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

model.evaluate(xtest, ytest)