Cannot export a slightly customized XLMRoberta model from keras_nlp

[YangIsNotAvailable]

Describe the bug Cannot export the model.

To Reproduce

import keras
from keras_nlp.models import XLMRobertaPreprocessor, XLMRobertaBackbone
import tensorflow as tf

preprocessor = XLMRobertaPreprocessor.from_preset("xlm_roberta_base_multi")
backbone = XLMRobertaBackbone.from_preset("xlm_roberta_base_multi")

inputs = keras.Input(shape=(), dtype=tf.string)
x = preprocessor(inputs)
x = backbone(x)
x = keras.layers.GlobalAveragePooling1D()(x)
outputs = keras.layers.Dense(10)(x)
model = keras.Model(inputs, outputs)

model.compile(optimizer=keras.optimizers.AdamW())

model.export("./test.tfsm")

Expected behavior Success.

Console output

AssertionError: Tried to export a function which references an 'untracked' resource. TensorFlow objects (e.g. tf.Variable) captured by functions must be 'tracked' by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly. See the information below:
        Function name = b'__inference_signature_wrapper___call___11987'
        Captured Tensor = <ResourceHandle(name="_0_SentencepieceOp", device="/job:localhost/replica:0/task:0/device:CPU:0", container="localhost", type="tensorflow::text::(anonymous namespace)::SentencepieceResource", dtype and shapes : "[  ]")>
        Trackable referencing this tensor = <tensorflow_text.python.ops.sentencepiece_tokenizer._SentencepieceModelResource object at 0x7fe279bef640>
        Internal Tensor = Tensor("11587:0", shape=(), dtype=resource)

[mehtamansi29]

Hi @YangIsNotAvailable

Thanks for reporting the issue. I have tested the code snippet and reproduces the reported behaviour. Attached gist file for reference.

We will look into the issue and update you the same.