Using tensorflow keras callbacks

my_callback.py

def train_model():
    # ... Define the model and training data etc here

    def my_callback(epoch, logs):
        # Optionally clear the file if it exists here? or use on_train_begin cb
        # Focus on the point here
        with open('train_logs.json', 'a') as l:
            data =json.dumps({'epoch': epoch, **logs})  
            l.write(data + '\n')

    my_callbacks = [
        tf.keras.callbacks.LambdaCallback(
            on_epoch_end=my_callback
            )
        ]

    batch_size = 512
    epochs = 10

    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

    model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1, callbacks=my_callbacks)

    return model