code: def test_data_generator_with_increasing_batch_size(
      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()
  
file_path: rasa/tests/utils/tensorflow/test_data_generator.py
imports: import pytest
  import scipy.sparse
  import numpy
  from rasa.utils.tensorflow.model_data import ragged_array_to_ndarray, FeatureArray, RasaModelData
  from rasa.utils.tensorflow.data_generator import RasaDataGenerator, RasaBatchDataGenerator
label_desc: Use image node_id 1 for calling a global function with example usage: test_data_generator_with_increasing_batch_size(model_data) without return types