class CustomCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        logging.info(f"Epoch {epoch} ended with loss: {logs['loss']} and accuracy: {logs['accuracy']}")

with strategy.scope():
    model = create_model()
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(train_dataset, epochs=5, validation_data=test_dataset,
              callbacks=[tensorboard_callback, CustomCallback()])

        __
        __