keras-nlp
keras-nlp copied to clipboard
Saving and loading TransformerDecoder/Encoder with custom objects rebuilds the layer.
Describe the bug When loading a model with TransformerDecoder, with the custom objects option, it rebuilds and reinitializes the layers within TransformerDecoder.
To Reproduce In keras_nlp/layers/transformer_decoder_test, change the test case of test_save_model by changing the load_model line to: (https://github.com/keras-team/keras-nlp/blob/38ca02b96b1a9100bcc76d123b399ecbc1f902c2/keras_nlp/layers/transformer_decoder_test.py#L194
loaded_model = keras.models.load_model(path, custom_objects={"TransformerDecoder": decoder})
After that, run the following command to run the test:
python3 -m pytest keras_nlp/layers/transformer_decoder_test.py
This will result in a mismatch error in test_save_model.
Expected behavior Layers within TransformerDecoder should not be reinitialized, the saved weights should be kept.