GPU memory abnormal increase

[duguwanglong]

Using the retvec model as embeding methods, but the GPU memory abnormal increase to 40G.
tensorflow 2.15.0
keras 2.15.0
tf-keras 2.15.0
retvec 1.0.1

Using methods:

def get_retvec_tokenizer(model_path):
    with tf.device('/CPU:0'):
        inputs = layers.Input(shape=(1,), dtype=tf.string)
        outputs = RETVecTokenizer(model=model_path)(inputs)
        retvec = tf.keras.Model(inputs=inputs, outputs=outputs)

    return retvec

Init the model state:

Using the model state inference: