Data Augmentation Multi Outputs

This question is asked several times here on SE, but I havent been able to find the right answer.
I’m trying to build a network with 1 input and 2 outputs. I don’t have a lot of data so I would like to use a generator for augmentation (preferably with imgaug).
My code:

seq = iaa.Sequential([ ....
])

gen = ImageDataGenerator(preprocessing_function=seq.augment_image)  
batch_size = 64

def generate_data_generator(generator, X, Y1, Y2):
    genX = gen.flow(X, batch_size=batch_size, seed=42)
    genY1 = gen.flow(Y1, batch_size=batch_size, seed=42)
    while True:
            Xi = genX.next()
            Yi1 = genY1.next()
            Yi2 = function(Y2)
            yield Xi, [Yi1, Yi2]

H = model.fit_generator(generate_data_generator(gen, trainX, trainY1, trainY2),
                epochs=epochs,
                steps_per_epoch=len(trainX) // batch_size,
                validation_data=(testX, [testY1, testY2]))

With this, I get the error:

'ValueError: ('Input data in `NumpyArrayIterator` should have rank 4. You passed an array with shape', (115, 16))' 

a normal fit like this works just fine, so I don’t there is anything wrong with the normal input

H = model.fit(trainX, {"output1": trainY1, "output2": trainY2},
    validation_data=(testX, {"output1": testY1, "output2": testY2}),
    epochs=epochs)

This is my input, I have no idea how I get a shape of (115,16) here

trainX.shape
trainY1[0].shape
trainY2[0].shape

Out[181]:
(115, 158, 100, 3)
Out[181]:
(16,)
Out[181]:
(9,)

thank you