Correct - and also from this thread Ray Iteration vs Keras Epoch, we recommend using: ray.air.integrations.keras.ReportCheckpointCallback — Ray 2.3.0, which works in both Tune training functions, as well as the AIR TensorflowTrainer.
Could you also mark this as resolved? Thanks!