code: def testRegisterFunction(self):
      if not common.has_eager_been_enabled():
          self.skipTest(
              "Only supported in TF2.x. Step is required in TF1.x"
          )
  
      time_step_spec = ts.TimeStep(
          step_type=tensor_spec.BoundedTensorSpec(
              dtype=tf.int32, shape=(), name="st", minimum=0, maximum=2
          ),
          reward=tensor_spec.BoundedTensorSpec(
              dtype=tf.float32,
              shape=(),
              name="reward",
              minimum=0.0,
              maximum=5.0,
          ),
          discount=tensor_spec.BoundedTensorSpec(
              dtype=tf.float32,
              shape=(),
              name="discount",
              minimum=0.0,
              maximum=1.0,
          ),
          observation=tensor_spec.BoundedTensorSpec(
              dtype=tf.float32,
              shape=(4,),
              name="obs",
              minimum=-10.0,
              maximum=10.0,
          ),
      )
      action_spec = tensor_spec.BoundedTensorSpec(
          dtype=tf.int32, shape=(), minimum=0, maximum=10, name="act_0"
      )
  
      network = q_network.QNetwork(
          input_tensor_spec=time_step_spec.observation,
          action_spec=action_spec,
      )
  
      policy = q_policy.QPolicy(
          time_step_spec=time_step_spec,
          action_spec=action_spec,
          q_network=network,
      )
  
      saver = policy_saver.PolicySaver(policy, batch_size=None)
      async_saver = async_policy_saver.AsyncPolicySaver(saver)
      async_saver.register_function(
          "q_network", network, time_step_spec.observation
      )
  
      path = os.path.join(self.get_temp_dir(), "save_model")
      async_saver.save(path)
      async_saver.flush()
      async_saver.close()
      self.assertFalse(async_saver._save_thread.is_alive())
      reloaded = tf.compat.v2.saved_model.load(path)
  
      sample_input = self.evaluate(
          tensor_spec.sample_spec_nest(
              time_step_spec.observation, outer_dims=(3,)
          )
      )
      expected_output, _ = network(sample_input)
      reloaded_output, _ = reloaded.q_network(sample_input)
  
      self.assertAllClose(expected_output, reloaded_output)
  
file_path: agents/tf_agents/policies/async_policy_saver_test.py
imports: import os
  from absl.testing.absltest import mock
  import tensorflow
  from tf_agents.networks import q_network
  from tf_agents.policies import async_policy_saver
  from tf_agents.policies import policy_saver
  from tf_agents.policies import q_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 6 for calling the AsyncPolicySaverTest obj's underlying member method code with example usage: obj.testRegisterFunction() without return types