I think a lot of these may have improved since your last experience with Keras. It's pretty easy to override the training loop and/or make custom loss. The belo

I think a lot of these may have improved since your last experience with Keras. It's pretty easy to override the training loop and/or make custom loss. The below is for overriding training / test step altogether, custom loss is easier by making a new loss function/class. https://keras.io/examples/keras_recipes/trainer_pattern/ https://keras.io/examples/keras_recipes/trainer_pattern/ > - Keras's training loop assumes you can fit all the data in memory and that the data is fully preprocessed, which in the world of LLMs and big data is infeasible. The Tensorflow backend has the excellent tf.data.Dataset API, which allows for out of core data and processing in a streaming way.