Custom train_step in subclass model not saved
Hi @rcalonso ,

Keras has introduced a decorator `@keras.saving.register_keras_serializable() ` which needs to be decorated on your custom class and save format should be `.keras`.

I have done the required changes and tested with Tf2.14 and keras-nightly(3.0.0.Dev) and both working fine and able to reload custom train_step properly.Please refer attached gists [Tf2.14](https://colab.sandbox.google.com/gist/SuryanarayanaY/38b7e326ca9139f59dcc81b4d055788b/62450-tf2-15.ipynb) and [keras3](https://colab.sandbox.google.com/gist/SuryanarayanaY/b78b698e83a39069e1a536c29836b6c8/62450-keras3-0-dev.ipynb).

Hope this resolves your query.

 Hi @SuryanarayanaY,

Many thanks! I tried before with the decorator `@keras.saving.register_keras_serializable()` but I was saving in `tf` format. Using the save format `.keras` works as expected.

Thanks for the help. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62450">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62450">No</a>

### Issue type

Bug

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

No

### Source

binary

### TensorFlow version

2.13.1 | 2.14.0 | 2.15.0

### Custom code

Yes

### OS platform and distribution

Google Colab

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Hi,

**Describe the current behavior**
I'm building a custom model via model subclassing following the steps in https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit. I'm overriding the `train_step` and `test_step` method. I can train the model without any problem but when I load a saved model, the `train_step` and `test_step` are not loaded and instead the methods from the super class `tf.keras.Model` are used.

I've tried the suggestions from previous issues where they had the same problem but the suggestions don't work. 

- https://github.com/tensorflow/tensorflow/issues/38103
- https://github.com/tensorflow/tensorflow/issues/46530

Also, I've gone through the documentation about serialization and saving https://www.tensorflow.org/guide/keras/serialization_and_saving#how_savedmodel_handles_custom_objects and the suggestion from https://github.com/tensorflow/tensorflow/issues/46530#issuecomment-790209731 but it doesn't work.

**Describe the expected behavior**
I'd expect that `train_step` and `test_step` are loaded from the saved model.

### Standalone code to reproduce the issue
https://gist.github.com/rcalonso/8b0c55f3e31b294a9fe68d0f1f277635


### Relevant log output

_No response_