Hi!
I have a functional model M that is the concatenation of 2 sub-models A and B:

class M(keras.Model)
def __init__(self, dataset):

self.a = A(...)
self.b = B(...)

self.build(dataset.get_input_shape())

def call(self, x):
x = self.a(x)
y = self.b(x)

return y

However, B is a multi-task head, so it returns a dictionary:

class B(keras.Model):
def __init__(self):
self.head1 = ....
def call(self, x):
return {
'head1': self.head1(x)
}

The issue is that when I call the summary on A or B, it works correctly (also with expand_nested=True); but when I use M.summary() it gives me the following error:

File "/home/matteo.doria/Desktop/FoundationalModel/.venv/lib/python3.12/site-packages/keras/src/utils/summary_utils.py", line 115, in format_layer_shape
return output_shapes[0]

Python: 3.12 Keras: 3.13.2 Backend: Tensorflow 2.20