I have defined a callback that runs on the epoch end and calculate the metrics. It is working fine in terms of calculating the desired metrics. Below is the function for reference

callback to find metrics at epoch end

class Metrics(tf.keras.callbacks.Callback):
    def __init__(self, train_tf_data, val_tf_data, model, CLASSES, logs={}, **kwargs):
        super().__init__(**kwargs)
        self.train_tf_data = train_tf_data
        self.val_tf_data = val_tf_data
        self.model = model
        self.CLASSES = CLASSES
        # for train data
        self.train_f1_after_epoch = 0
        self.train_prec_after_epoch = 0
        self.train_recall_after_epoch = 0
        # for val data
        self.val_f1_after_epoch = 0
        self.val_prec_after_epoch = 0
        self.val_recall_after_epoch = 0

    def on_train_begin(self, logs={}):
        self.train_reports = None
        self.val_reports = None
        self.val_f1_after_epoch = 0

    def on_epoch_end(self, epoch, logs={}):
        # for train data
        self.train_reports = test_model(model=self.model, data=self.train_tf_data, 
                                        CLASSES=self.CLASSES)
        self.train_f1_after_epoch = self.train_reports['f1_score']
        self.train_recall_after_epoch = self.train_reports['recall']
        self.train_prec_after_epoch = self.train_reports['precision']

        # for val data
        self.val_reports = test_model(model=self.model, data=self.val_tf_data, 
                                      CLASSES=self.CLASSES)
        self.val_f1_after_epoch = self.val_reports['f1_score']
        self.val_recall_after_epoch = self.val_reports['recall']
        self.val_prec_after_epoch = self.val_reports['precision']

        # saving train results to log dir
        logs["train_f1_after_epoch"]=self.train_f1_after_epoch
        logs['train_precision_after_epoch'] = self.train_prec_after_epoch
        logs['train_recall_after_epoch'] = self.train_recall_after_epoch
        
        # saving val results to log dir
        logs['val_f1_after_epoch'] = self.val_f1_after_epoch
        logs['val_precision_after_epoch'] = self.val_prec_after_epoch
        logs['val_recall_after_epoch'] = self.val_recall_after_epoch


        print('train_reports_after_epoch', self.train_reports)
        print('val_reports_after_epoch', self.val_reports)

Code for test_model

def test_model(model, data, CLASSES, label_one_hot=True, average="micro"):
    images_ds = data.map(lambda image, label: image)
    labels_ds = data.map(lambda image, label: label).unbatch()
    NUM_VALIDATION_IMAGES = count_data_items(tf_records_filenames=data)
    cm_correct_labels = next(iter(labels_ds.batch(NUM_VALIDATION_IMAGES))).numpy() # get everything as one batch
    if label_one_hot is True:
        cm_correct_labels = np.argmax(cm_correct_labels, axis=-1)
    cm_probabilities = model.predict(images_ds)
    cm_predictions = np.argmax(cm_probabilities, axis=-1)
    
    # cmat = confusion_matrix(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)))

    warnings.filterwarnings('ignore')
    score = f1_score(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)), average=average)
    precision = precision_score(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)), average=average)
    recall = recall_score(cm_correct_labels, cm_predictions, labels=range(len(CLASSES)), average=average)
    # cmat = (cmat.T / cmat.sum(axis=1)).T # normalized
    # print('f1 score: {:.3f}, precision: {:.3f}, recall: {:.3f}'.format(score, precision, recall))
    test_results = {'f1_score': score, 'precision':precision, 'recall':recall}
    warnings.filterwarnings('always')
    return test_results

Some model code.....

Model code

m1 = tf.keras.metrics.CategoricalAccuracy()
m2 = tf.keras.metrics.Recall()
m3 = tf.keras.metrics.Precision()
m4 = Metrics(train_tf_data=train_data, 
             val_tf_data=test_data, model=model, 
             CLASSES=CLASS_NAMES)
optimizers = [
        tfa.optimizers.AdamW(learning_rate=lr * .001 , weight_decay=wd),
        tfa.optimizers.AdamW(learning_rate=lr, weight_decay=wd)

           ]
optimizers_and_layers = [(optimizers[0], model.layers[0]), (optimizers[1], model.layers[1:])]
    
optimizer = tfa.optimizers.MultiOptimizer(optimizers_and_layers)


model.compile(
    optimizer= optimizer,
    loss = 'categorical_crossentropy',
    metrics=[m1, m2, m3],
    )

Using this in the callback

checkpoint_cb = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path, 
                                                    monitor = 'val_f1_after_epoch',
                                                    save_best_only=True,
                                                    save_weights_only=True,
                                                    mode='max',
                                                    save_freq='epoch',
                                                    verbose=1)
                                                    
checkpoint_cb._supports_tf_logs = False

The issue that I am facing is that it is giving me a warning that say

WARNING:tensorflow:Can save best model only with val_f1_after_epoch available, skipping

Upon investigating history I found that metrics is available in the history

print(list(history.history.keys()))
['loss',
'categorical_accuracy',
'recall',
'precision',
'val_loss',
'val_categorical_accuracy',
'val_recall',
'val_precision',
'train_f1_after_epoch',
'train_precision_after_epoch',
'train_recall_after_epoch',
'val_f1_after_epoch', #this is the metrics
'val_precision_after_epoch',
'val_recall_after_epoch']

Please let me know what I am missing here, I want to save the best model based on my custom metrics?