model_subclass.py

class MNIST(keras.Model):
    def __init__(self):
        super().__init__()
        self.conv_1 = layers.Conv2D(32, kernel_size=(3, 3), activation="relu")
        self.conv_2 = layers.Conv2D(64, kernel_size=(3, 3), activation="relu")
        self.max_pool = layers.MaxPooling2D(pool_size=(2, 2))
        self.flatten = layers.Flatten()
        self.dropout = layers.Dropout(0.5)
        self.out =  layers.Dense(num_classes, activation="softmax")

    def call(self, inputs):
        x = self.conv_1(inputs)
        x = self.max_pool(x)
        x = self.conv_2(x)
        x = self.max_pool(x)
        x = self.flatten(x)
        x = self.dropout(x)
        return self.out(x)

model = MNIST()