KerasRegressor with multiple outputs can be trained, scores, but does not predict

Intro

I'm using the KerasRegressor wrapper for scikit-learn and the Functional API to create a model with multiple outputs. My model has two branches: one that predicts 1 value, another one that predicts 3 values at the same time. The keras version used is the latest one available (2.3.1 on win64)

Issue

The problem I encountered is the following: my model can be trained through the pipeline, I can get a score, but I cannot predict.

# To fit the pipeline, this line runs successfully
pipeline.fit(X_train, [y_train_branch_1, y_train_branch_2])`

# To get the score, it works as well
pipeline.score(X_test, [y_test_branch_1, y_test_branch_2])

# To make predictions however, it doesn't
pipeline.predict(X_test)

The last line does not work:

ValueError: could not broadcast input array from shape (11963,3) into shape (11963)

It happens here:

~\Anaconda3\envs\blades\lib\site-packages\keras\wrappers\scikit_learn.py in predict(self, x, **kwargs)
    320         """
    321         kwargs = self.filter_sk_params(Sequential.predict, kwargs)
--> 322         preds = np.array(self.model.predict(x, **kwargs))
    323         if preds.shape[-1] == 1:
    324             return np.squeeze(preds, axis=-1)

The input shapes are:

• X_train:(47849, 4)
• y_train_branch_1: (47849, 3)
• y_train_branch_2: (47849, 1)
• X_test: (11963, 4)
• y_test_branch_1: (11963, 3)
• y_test_branch_2: (11963, 1)

It seems that the results of the predictions, which come in two parts, are not properly merged when using .predict. Has anyone ran into this issue before ? Concatenating the two outputs in the model directly may be a way to work around, but I would like to solve this issue anyway.

Thanks for your help!

Code

Imports

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from keras.wrappers.scikit_learn import KerasRegressor
from keras.callbacks import EarlyStopping
from keras.callbacks import ModelCheckpoint
from keras import Input, Model
from keras.layers import Dense, Activation, Concatenate
from keras.optimizers import Adam

Model definition

def build_branch_1(inputs):
    x = Dense(12)(inputs)
    x = Activation('elu')(x)
    x = Dense(12)(x)
    x = Activation('elu')(x)
    
    a = Dense(12)(x)
    a = Activation('elu')(a)
    a = Dense(1)(a)
    b = Dense(12)(x)
    b = Activation('elu')(b)
    b = Dense(1)(b)
    c = Dense(12)(x)
    c = Activation('elu')(c)
    c = Dense(1)(c)
    
    x = Concatenate(name='branch_1')([a, b, c])
    return x
    
def build_branch_2(inputs):
    x = Dense(12)(inputs)
    x = Activation('elu')(x)
    x = Dense(12)(x)
    x = Activation('elu')(x)
    x = Dense(12)(x)
    x = Activation('elu')(x)
    x = Dense(1, name='branch_2')(x)
    return x

def build_model():
    inputs = Input(shape=(4,))
    branch_1 = build_branch_1(inputs)
    branch_2 = build_branch_2(inputs)
    
    model = Model(inputs=inputs, outputs=[branch_1, branch_2])
    model.compile(optimizer=Adam(), loss='mse')
    return model

Pipeline definition

pipeline = Pipeline([
    ('stdscaling', StandardScaler()),
    ('model', KerasRegressor(
        build_fn=build_model, 
        batch_size=128,
        epochs=10,
        verbose=2,
        validation_split=.2,
        callbacks=[
            EarlyStopping(monitor='loss', min_delta=.01, patience=250, 
                          restore_best_weights=True),
            ModelCheckpoint(filepath='model/model_' + model_name + '.hdf5', 
                            monitor='val_loss', 
                            save_best_only=True,
                            save_weights_only=False)
        ]
    ))
])