X = layers.Flatten()(last_output)

Have a look on this post: Python Calling