class MyModelHander(ModelHandler): ... def batch_elements_kwargs(self) -> Dict[str, Any]: return { 'min_batch_size': 1, 'max_batch_size': 5, }