Tensorflow 2.2: Dataset with increasing batch size

I have a custom model to train a text classifier. In our code we overwrote all important methods of tf.keras.models.Model, e.g. fit, predict etc. One of the reasons for that was that we wanted to adapt the batch size every epoch. So, during the first epoch we would have, for example, a batch size of 16, in the next iteration the batch size would increase to 20, and so on. Until we reach an upper limit.

Here is the method that generates our tf.Dataset. We would call this method for every epoch to get a dataset with the desired batch size.

def as_tf_dataset(
    self, batch_size: int, batch_strategy: Text = SEQUENCE, shuffle: bool = False
) -> tf.data.Dataset:
    """Create tf dataset."""

    shapes, types = self._get_shapes_types()

    return tf.data.Dataset.from_generator(
        lambda batch_size_: self._gen_batch(batch_size_, batch_strategy, shuffle),
        output_types=types,
        output_shapes=shapes,
        args=([batch_size]),
    )

With Tensorflow 2.2. it is possible to just overwrite train_step and test_step instead of fit etc. This would simplify our code quite a bit. However, I could not find a way how to keep the increasing batch size over the epochs.

Does someone has an idea on how to solve this? Or do we need to keep our custom fit method in order to achieve this?