[BUG] when the model is loaded and saved again, the argument names disappear (become 'args_0, args_0_1 etc)

[rnyak]

Bug description

when the model (it is transformer based XLNet model that we are training) is loaded and saved again, the argument names disappear (become 'args_0, args_0_1 etc) and the saved model becomes invalid. Because of that we cannot export the reloaded model using Merlin Systems, therefore cannot generate the config files to serve on Triton.

Steps/Code to reproduce bug

Please run these two gists in order to repro the issue:

1. https://gist.github.com/rnyak/0baa9e80cb419128379f5e91e8a30ed5
2. https://gist.github.com/rnyak/d222fe2eb7b2ecbf87d924a762c1f7d4#file-sbr_xlnet_load-py

Expected behavior

Environment details

• Merlin version: We are using merlin-tensorflow:23.08 image with the latest main branch pulled and installed.
• Platform:
• Python version:
• PyTorch version (GPU?):
• Tensorflow version (GPU?):

[rnyak]

One quick solution to that is to add model block line right before loading the saved model. Something like that:

_ =  mm.XLNetBlock(d_model=32, n_head=2, n_layer=2)  # replace XLNetBlock with `GPT2Block` if you use GPT2 block in model training.
loaded_model = tf.keras.models.load_model('/workspace/saved_model/')