Calling submodels in train_step throws ValueError when saving and loading error.
@hubgrim I am able to get a different error while reproducing the issue reported here. Kindly check the [gist](https://colab.research.google.com/gist/sushreebarsa/8965a0e01764c86c092f2a2d6d86959f/64576.ipynb) for reference. 
The .keras extension typically represents the recommended Keras v3 format. If your model was saved in a different format (e.g., HDF5 - .h5), you'll need to use the appropriate loading function as follows;
```
loaded_model = tf.keras.models.load_model('custom_model.h5')  # For HDF5 format

```
Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. > @hubgrim I am able to get a different error while reproducing the issue reported here. Kindly check the [gist](https://colab.research.google.com/gist/sushreebarsa/8965a0e01764c86c092f2a2d6d86959f/64576.ipynb) for reference. The .keras extension typically represents the recommended Keras v3 format. If your model was saved in a different format (e.g., HDF5 - .h5), you'll need to use the appropriate loading function as follows;
> 
> ```
> loaded_model = tf.keras.models.load_model('custom_model.h5')  # For HDF5 format
> ```
> 
> Thank you!

Hm, I don't know why this error differs or why the saving format would be an issue, since I use the .keras extension when loading and saving. The lines are right next to each other. @hubgrim While the .keras extension suggests a Keras model, there may be compatibility issues between the version of Keras you used to save the model and the version you're using to load it.  Even small version bumps can sometimes cause errors. Please ensure the version compatibility.
Thank you! The error arises even when saving/loading within the same script. I did not save the model at some point and reloaded it at some other time, but within the same call of the program. Both the model, that provides .save(), and load_model() were imported from tensorflow.keras.models so I don't see how this could be rooted in version incompatibility.
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.15

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 23.10

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

When saving and loading a nested model I got this error.
```ValueError: Layer 'dense_3' expected 0 variables, but received 2 variables during loading. Expected: []```
Which was weird, because there was no dense_3 layer. 

Calling the submodels by themselves over self.net1 and self.net2 **inside** the train_step would throw the error, but having them called by the call method of the higher model and returning their respective values does not throw an error.
How this could lead to this error messages is beyond me.
I managed to build the following minimal example:

### Standalone code to reproduce the issue

```shell
keras
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Activation, Input, BatchNormalization, Dropout, Flatten, Identity


args = {"activation": "relu",
        "batch_norm": True}

@keras.saving.register_keras_serializable()
class CustomModel1(Model):
    def __init__(self):
        super().__init__()
        self.dense = Dense(32)

    def call(self, inputs):
        x = self.dense(inputs)
        return x


@keras.saving.register_keras_serializable()
class CustomModel2(Model):
    def __init__(self):
        super().__init__()
        self.dense = Dense(32)

    def call(self, inputs):
        x = self.dense(inputs)
        return x


@keras.saving.register_keras_serializable()
class CustomModel3(Model):
    def __init__(self):
        super().__init__()
        self.net1 = CustomModel1()
        self.net2 = CustomModel2()

    def call(self, inputs):
        z = self.net1(inputs)
        x = self.net2(z)
        return z, x

    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            # z, y_pred = self(x)                 # this fixes it instead
            y_pred = self.net2(self.net1(x))      # this line throws the error
            loss = self.compiled_loss(y, y_pred)

        trainable_vars = self.trainable_weights
        gradients = tape.gradient(loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        self.compiled_metrics.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}


# Instantiate the model
model = CustomModel3()

# Compile the model
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Create some dummy data for training
x_train = np.random.random((1000, 32))
y_train = np.random.randint(10, size=(1000,))

# Train the model for one epoch
model.fit(x_train, y_train, epochs=1)

# Save the model
model.save('custom_model.keras', save_format='keras')

# Load the model again
loaded_model = tf.keras.models.load_model('custom_model.keras')

# Generate some sample data for prediction
x_sample = np.random.random((10, 32))  # Assuming 10 samples with 32 features each

# Make predictions using the loaded model
predictions = loaded_model.predict(x_sample)
print(predictions)
# Print the predictions
print(model.summary())
```


### Relevant log output

```shell
/home/georg/anaconda3/envs/vae/bin/python3.11 /home/georg/Python Projects/vae/tests/test_saving_basic.py 
2024-03-27 10:02:21.156169: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE4.1 SSE4.2 AVX AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
32/32 [==============================] - 1s 2ms/step - loss: 9.2085 - accuracy: 0.0400
/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_api.py:164: UserWarning: You are saving a model that has not yet been built. It might not contain any weights yet. Consider building the model first by calling it on some data.
  saving_lib.save_model(model, filepath)
WARNING:absl:Skipping variable loading for optimizer 'Adam', because it has 1 variables whereas the saved optimizer has 9 variables. 
Traceback (most recent call last):
  File "/home/georg/Python Projects/vae/tests/test_saving_basic.py", line 79, in <module>
    loaded_model = tf.keras.models.load_model('custom_model.keras')
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_api.py", line 254, in load_model
    return saving_lib.load_model(
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_lib.py", line 281, in load_model
    raise e
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_lib.py", line 269, in load_model
    _load_state(
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_lib.py", line 457, in _load_state
    _load_state(
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_lib.py", line 466, in _load_state
    _load_container_state(
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_lib.py", line 534, in _load_container_state
    _load_state(
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/saving/saving_lib.py", line 435, in _load_state
    trackable.load_own_variables(weights_store.get(inner_path))
  File "/home/georg/anaconda3/envs/vae/lib/python3.11/site-packages/keras/src/engine/base_layer.py", line 3531, in load_own_variables
    raise ValueError(
ValueError: Layer 'dense_3' expected 0 variables, but received 2 variables during loading. Expected: []

Process finished with exit code 1
```
