Call _set_save_spec() when creating TF models
Much like confused ducklings, subclassed Keras models tend to imprint on the first concrete input shapes they see unless we explicitly `build()` them with more general shapes. However, the default `tf.keras.Model.build()` makes several restrictive assumptions and doesn't work for us.

The solution is to directly call `model._set_save_spec()` with the shapes we want. We do this in the `__init__` to make sure that it happens before the model is built or called with any inputs. We can also now remove the override on `model.save()`, which is no longer necessary now that we're fixing this properly.

Fixes #19231
_The documentation is not available anymore as the PR was closed or merged._