code: def save_test(alg_name, framework="tf", multi_agent=False):
      config = (
          get_trainable_cls(alg_name)
          .get_default_config()
          .framework(framework)
          # Switch on saving native DL-framework (tf, torch) model files.
          .checkpointing(export_native_model_files=True)
      )
  
      if alg_name in RLMODULE_SUPPORTED_ALGOS:
          config = config.experimental(_enable_new_api_stack=False)
  
      if "DDPG" in alg_name or "SAC" in alg_name:
          config.environment("Pendulum-v1")
          algo = config.build()
          test_obs = np.array([[0.1, 0.2, 0.3]])
      else:
          if multi_agent:
              config.multi_agent(
                  policies={"pol1", "pol2"},
                  policy_mapping_fn=(
                      lambda agent_id, episode, worker, **kwargs: "pol1"
                      if agent_id == "agent1"
                      else "pol2"
                  ),
              )
              config.environment(
                  MultiAgentCartPole, env_config={"num_agents": 2}
              )
          else:
              config.environment("CartPole-v1")
          algo = config.build()
          test_obs = np.array([[0.1, 0.2, 0.3, 0.4]])
  
      export_dir = os.path.join(
          ray._private.utils.get_user_temp_dir(),
          "export_dir_%s" % alg_name,
      )
  
      algo.train()
      print("Exporting algo checkpoint", alg_name, export_dir)
      export_dir = algo.save(export_dir).checkpoint.path
      model_dir = os.path.join(
          export_dir,
          "policies",
          "pol1" if multi_agent else DEFAULT_POLICY_ID,
          "model",
      )
  
      # Test loading exported model and perform forward pass.
      if framework == "torch":
          filename = os.path.join(model_dir, "model.pt")
          model = torch.load(filename)
          assert model
          results = model(
              input_dict={"obs": torch.from_numpy(test_obs)},
              # TODO (sven): Make non-RNN models NOT expect these args at all.
              state=[torch.tensor(0)],  # dummy value
              seq_lens=torch.tensor(0),  # dummy value
          )
          assert len(results) == 2
          assert results[0].shape == (1, 2)
          assert results[1] == [torch.tensor(0)]  # dummy
      else:
          model = tf.saved_model.load(model_dir)
          assert model
          results = model(
              tf.convert_to_tensor(test_obs, dtype=tf.float32)
          )
          assert len(results) == 2
          assert results[0].shape == (1, 2)
          # TODO (sven): Make non-RNN models NOT return states (empty list).
          assert results[1].shape == (1, 1)  # dummy state-out
  
      shutil.rmtree(export_dir)
  
file_path: ray/rllib/algorithms/tests/test_algorithm_export_checkpoint.py
imports: import numpy
  import os
  import shutil
  import unittest
  import ray
  from ray.rllib.examples.env.multi_agent import MultiAgentCartPole
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.test_utils import framework_iterator
  from ray.tune.registry import get_trainable_cls
label_desc: Use image node_id 1 for calling a global function with example usage: save_test(alg_name, framework, multi_agent) without return types