def METHOD_NAME(model_data: RasaModelData):
    epochs = 2
    data_generator = RasaBatchDataGenerator(
        model_data,
        batch_size=[1, 2],
        epochs=epochs,
        batch_strategy="balanced",
        shuffle=True,
    )
    expected_batch_sizes = [[1, 1, 1, 1, 1], [2, 2, 1]]
    for _epoch in range(epochs):
        iterator = iter(data_generator)
        assert len(data_generator) == len(expected_batch_sizes[_epoch])
        for i in range(len(data_generator)):
            batch, _ = next(iterator)
            assert len(batch) == 11
            assert len(batch[0]) == expected_batch_sizes[_epoch][i]
        with pytest.raises(StopIteration):
            next(iterator)
        data_generator.on_epoch_end()