Wrapping custom model subclassing tf.keras.Model does not work

[vboussange]

Hey there,
I have tried to use scikeras with a custom class inheriting keras.Model, but unfortunately this seems to fail.

Here is MWE:

from sklearn.preprocessing import FunctionTransformer
from scikeras.wrappers import KerasRegressor
from typing import Dict, Iterable, Any
from scikeras.wrappers import KerasRegressor
import tensorflow as tf
from tensorflow import keras
from keras.layers import Dense
import numpy as np

class MyCustomModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.nn1 = Dense(1)

    def call(self,inputs):
        return self.nn1(inputs)

class MyWrapper(KerasRegressor):
    def __init__(
        self,
        optimizer="adam",
        optimizer__learning_rate=0.001,
        epochs=200,
        verbose=0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.optimizer = optimizer
        self.epochs = epochs
        self.verbose = verbose

    def _keras_build_fn(self, compile_kwargs: Dict[str, Any]):
        model = MyCustomModel()
        model.compile(loss="mse", optimizer=compile_kwargs["optimizer"])
        return model
    
    
X = np.random.random(size=(100, 1))
y = np.sum(X, axis=1)
model = MultiInputModel()
model.fit(X,y)

which throws

File [/opt/homebrew/Caskroom/mambaforge/base/envs/SEAM/lib/python3.11/site-packages/scikeras/wrappers.py:549](https://file+.vscode-resource.vscode-cdn.net/opt/homebrew/Caskroom/mambaforge/base/envs/SEAM/lib/python3.11/site-packages/scikeras/wrappers.py:549), in BaseWrapper._check_model_compatibility(self, y)
    544 # check if this is a multi-output model
    545 if getattr(self, "n_outputs_expected_", None):
    546     # n_outputs_expected_ is generated by data transformers
    547     # we recognize the attribute but do not force it to be
    548     # generated
--> 549     if self.n_outputs_expected_ != len(self.model_.outputs):
    550         raise ValueError(
    551             "Detected a Keras model input of size"
    552             f" {self.n_outputs_expected_ }, but {self.model_} has"
    553             f" {len(self.model_.outputs)} outputs"
    554         )
    555 # check that if the user gave us a loss function it ended up in
    556 # the actual model

TypeError: object of type 'NoneType' has no len()

Any thoughts on how to fix this? It seems that models instantiated with the functional API (keras.Model(inputs, outputs)) do not have the same outputs attribute as those created from user-defined classes which inherit from keras.Model, which possibly explains this error.