Errors when trying to save and load custom Tensorflow model

Question

I want to save a custom model , i inspire the solution from https://www.tensorflow.org/guide/saved_model?hl=fr#specifying_signatures_during_export , i have a model class that inherit from tf.Module and i put the decorator @tf.function in the function that i want to call later after loading the model but i got an error . i use:

• tensorflow==2.5.0
• keras-rl2==1.0.5

First Method

from rl.agents.dqn import DQNAgent

class CustomModel(tf.Module):
     
  ......

   @tf.function(input_signature=[tf.TensorSpec(shape=(None, 94), dtype=tf.float32)])
 
   def predict(self, s):
        q_values=self.agent.model.predict(s[:,np.newaxis,:],steps=1)
        scores=q_values[:,1]
        return scores

I got an error "RuntimeError: Cannot get session inside Tensorflow graph function

Second Method:

    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 94), dtype=tf.float32)])
    def get_label(self,x):
       
        labels=self.predict_label(x)
     
        return labels


    def predict_label(self,s):
      

        # Define a custom predict function without using self.agent.model.predict
        def custom_predict(x):
            # Convert the tensor x to a NumPy array using tf.numpy_function
            def convert_to_numpy(y):
                return y.numpy()

            x_array = tf.numpy_function(convert_to_numpy, [x], tf.float32)

            # Perform prediction using the model on the NumPy array
            q_values = self.predict_with_model(x_array)

            return q_values

        # Use tf.numpy_function to execute the custom_predict function
        q_values = tf.numpy_function(custom_predict, [s], tf.float32)

    
      
        return q_values

    def predict_with_model(self, x):
        # Perform the prediction using self.agent.model.predict
        q_values = self.agent.model.predict(x[:, np.newaxis, :], steps=1)
        return q_values

the custom is saved without errors but when i load the model :

model= tf.saved_model.load("model_saved")
print(dir(model))
q_values=model.get_label(test_X)

i got an error "ValueError: callback pyfunc_0 is not found

Can you help me to understand the problem and be able to save and load my model with the predict function ?