Convert tensorflow saved model from 2.3 to >=2.5 to use tf.GradientTape(persistent=True)

TensorFlow General Discussion
, ,
1

Problem: I converted a tf1 model from this repo to tf2 and saved it as a SavedModel using tensorflow 2.3. As highlighted by this comment, when saving a SavedModel with a tf version anterior to 2.5, the model cannot be used with tf.GradientTape(persistent=True). However, this is exactly what I need to do. Is there a way to make my model usable with tf.GradientTape(persistent=True)? PS: I lost all the code I wrote to do the conversion.

Code example:
The model I am using can be downloaded from google drive here.

import tensorflow as tf
import tensorflow_addons as tfa

MODEL_FOLDER_PATH = ""
model = tf.saved_model.load(MODEL_FOLDER_PATH, tags="serve").signatures[
    "serving_default"
]
tfa.register_all(custom_kernels=False)  # required since tensorflow addons is used in the original model

@tf.function
def without_persistent_gradient(model):
    with tf.GradientTape():
        model(tf.ones((1, 64, 256, 3)))

@tf.function
def with_persistent_gradient(model):
    with tf.GradientTape(persistent=True):
        model(tf.ones((1, 64, 256, 3)))

without_persistent_gradient(model)   # This works!
with_persistent_gradient(model)  # This raises an error!

The error obtained is the following:

ValueError: Internal error: Tried to take gradients (or similar) of a variable without handle data:
Tensor("Backward/Predictor/decoder/while:13", shape=(), dtype=resource)

What I tried:
I tried to resave the model using tf2.8:

tf.saved_model.save(model, "my/path")

which raises the following error:

KeyError: "Failed to add concrete function 'b'__inference_pruned_12028'' to object-based SavedModel as it captures tensor <tf.Tensor: shape=(), dtype=resource, value=<Resource Tensor>> which is unsupported or not reachable from root. One reason could be that a stateful object or a variable that the function depends on is not assigned to an attribute of the serialized trackable object (see SaveTest.test_captures_unreachable_variable)."

So converted it to a Keras Model using the tensorflow_hub lib, but I am still getting the error with tf.GradientTape(persistent=True).

from tensorflow_hub import KerasLayer
input_layer = tf.keras.layers.Input((64, 256, 3))
model_layer = KerasLayer("my/path", tags="serve",signature="serving_default", signature_outputs_as_dict=True)
model = tf.keras.Model(input_layer, model_layer(input_layer))

Specs

python 3.9.12; linux 18.04; CPU

tensorflow==2.8; tensorflow-addons==0.16