class MyCustomModel(keras.Model):
    def __init__(self, units=32, num_classes=10):
        super(MyCustomModel, self).__init__()
        self.dense1 = MyCustomLayer(units, activation='relu')
        self.dense2 = MyCustomLayer(num_classes, activation='softmax')

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

model = MyCustomModel(units=64, num_classes=10)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

        __
        __