complex data generator for multitask pipe

multitask_data_generator.py

class MultitaskDataGenerator:
    def __init__(self, generators):
        self.generators = generators
        self.generator = self.generate_batch()
        
    def generate_batch(self, batch_size=None):
        while True:
            batch = self.__next__()
            yield batch
        
    def __next__(self):
        data_arrays = []
        for gen in self.generators:
            gen_data, gen_labels = next(gen.generator)

            if type(gen_data) not in {list, tuple}:
                gen_data = [gen_data]

            data_arrays += gen_data

        return data_arrays, gen_labels
            
    def __iter__(self):
        return self.__next__()