iris-classifier.py

class IrisClassifier(Model):
  def __init__(self):
    super(IrisClassifier, self).__init__()
    self.layer1 = Dense(10, activation='relu')
    self.layer2 = Dense(10, activation='relu')
    self.outputLayer = Dense(3, activation='softmax')

  def call(self, x):
    x = self.layer1(x)
    x = self.layer2(x)
    return self.outputLayer(x)
  
model = IrisClassifier()

model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss='categorical_crossentropy',
              metrics=['accuracy'])