[QUESTION]In 16_nlp_with_rnns_and_attentions.jpynb, when using manual masking in code block(44-45)

[iCurious-D]

In 16_nlp_with_rnns_and_attentions.jpynb,
when using manual masking in code block(44-45),

inputs = tf.keras.layers.Input(shape=[], dtype=tf.string)
token_ids = text_vec_layer(inputs)
mask = tf.math.not_equal(token_ids, 0)
Z = tf.keras.layers.Embedding(vocab_size, embed_size)(token_ids)
Z = tf.keras.layers.GRU(128, dropout=0.2)(Z, mask=mask)
outputs = tf.keras.layers.Dense(1, activation="sigmoid")(Z)
model = tf.keras.Model(inputs=[inputs], outputs=[outputs])

IDE raises a ValueError: A KerasTensor cannot be used as input to a TensorFlow function.
I guess this is because of my keras version?
Anyway, IDE suggests me to wrap tf_fn in a layer, and it worked!

inputs = keras.layers.Input(shape=[], dtype=tf.string)
class MyLayer(keras.layers.Layer):
    def call(self, x):
        token_ids = text_vec_layer(x)
        mask = tf.math.not_equal(token_ids, 0)
        return token_ids, mask
token_ids, mask = MyLayer()(inputs)
Z = keras.layers.Embedding(vocab_size, embed_size)(token_ids)
Z = keras.layers.GRU(128, dropout=0.2)(Z, mask=mask)
outputs = keras.layers.Dense(1, activation="sigmoid")(Z)
model = keras.Model(inputs=[inputs], outputs=[outputs])

My env_lib versions are:
keras 3.6.0
tensorflow 2.18.0

Thought it can work, I am still not sure is it possible to throw the wrapper away(i.e. can we directly use?).
And hopefully this isuue can help someone who encountered similar problems.