Insufficient documentation for multi-label training for numpy arrays

[PrasadNR]

I am working on medical imaging research and I am supposed to have one input and two labels (two heads/two outputs).

Reproducing the issue from https://datascience.stackexchange.com/questions/74911/data-augmentation-for-multiple-output-heads-in-keras (my own question of that link):

I have a transfer learning based two output classification problem. So, accordingly, I have formatted my data to have X_train as a (number of samples, height, width, channels) numpy array, y_train1 as (number of samples,) numpy array and y_train2 as (number of samples,) numpy array.

As I am not training using directory structure, I am using ImageDataGenerator.flow(). However, I am not able to figure out how I can pass two label arrays because, it is taking the labels as (2, number of samples) when I send it as [y_train1, y_train2] list.

I am able to train the network without Keras data augmentation (for two outputs). But, I am not able to apply data augmentation.

I am trying to do the following:

datagen = ImageDataGenerator(horizontal_flip=True, vertical_flip=True, rescale=1./255, class_mode="multi-label")
    
model.fit(datagen.flow(X_train, [y_train1, y_train2], batch_size=batch_size), batch_size=batch_size, epochs=nb_epochs, steps_per_epoch=spe, validation_data=(X_val, [y_val1, y_val2]))

Also, [ImageDataGenerator.flow][1] does not have class_mode unlike ImageDataGenerator.flow_from_dataframe.

I am raising this as an official issue as I have not been able to find any documentation on writing custom data augmentation class if the default ImageDataGenerator class cannot handle multiple labels. I hope I don't have to switch back to PyTorch because, Keras is what has gotten me to where I am in my career.