def _my_model_fn(features, labels, mode): my_head = tf.estimator.MultiClassHead(n_classes=3) logits = tf.keras.Model(...)(features)