TF-Agent have ValueError: Could not find matching function to call loaded from the SavedModel. Not dtype error

Question

I have trained a DQN agent and have some problems when I'm in action with time_step.

But I think the save and load policy is working a good job.

# Save Policy
my_policy       = tf_agent.collect_policy
saver           = policy_saver.PolicySaver(my_policy, batch_size=None)
saver.save(AGENT_MODEL_PATH)

# Load Policy
policy          = tf.compat.v2.saved_model.load(AGENT_MODEL_PATH)
policy_state    = policy.get_initial_state(batch_size=3)

Now I want to predict the action based on the provided timestep.

# Live
live_env            = Live_Environment( some parameter )
live_env            = tf_py_environment.TFPyEnvironment(live_env)

time_step           = live_env.current_time_step()

def run_step(step, sc):

    action_step     = policy.action(step, policy_state)
    next_time_step  = live_env.step(action_step.action)
    time_step       = next_time_step

    print(action_step.action)
    print(time_step.observation)
    print(time_step.reward)

    s.enter(10, 1, run_step, (step, sc))

s.enter(10, 1, run_step, (time_step, s ))
s.run()

But it's not working and I am getting the error :(

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (2 total):
    * 
    TimeStep(
    step_type   = <tf.Tensor 'time_step:0' shape=(1,) dtype=int32>, 
    reward      = <tf.Tensor 'time_step_1:0' shape=(1,) dtype=float32>, 
    discount    = <tf.Tensor 'time_step_2:0' shape=(1,) dtype=float32>, 
    observation = <tf.Tensor 'time_step_3:0' shape=(1, 26) dtype=float32>
    )
    * ()
  Keyword arguments: {}

Expected these arguments to match one of the following 2 option(s):

Option 1:
  Positional arguments (2 total):
    * 
    TimeStep(
    step_type   = TensorSpec(shape=(None,), dtype=tf.int32, name='step_type'), 
    reward      = TensorSpec(shape=(None,), dtype=tf.float32, name='reward'), 
    discount    = TensorSpec(shape=(None,), dtype=tf.float32, name='discount'), 
    observation = TensorSpec(shape=(None, 61), dtype=tf.float32, name='observation')
    )
    * ()
  Keyword arguments: {}

Option 2:
  Positional arguments (2 total):
    * 
    TimeStep(
    step_type   = TensorSpec(shape=(None,), dtype=tf.int32, name='time_step/step_type'), 
    reward      = TensorSpec(shape=(None,), dtype=tf.float32, name='time_step/reward'), 
    discount    = TensorSpec(shape=(None,), dtype=tf.float32, name='time_step/discount'), 
    observation = TensorSpec(shape=(None, 61), dtype=tf.float32, name='time_step/observation')
    )
    * ()
  Keyword arguments: {}

I search for 4 days but I can't find a solution. I am not sure why TimeStep is not found a matching function.

Is it because of the observation shape?

This is my first question, please be gentle.