code: def test_save_custom_objects(request, model_class):
      """
      Test feeding custom_objects in model.load(...) function
      """
      # Skip DQN (not an actor-critic policy)
      if model_class == DQN:
          return
  
      model_fname = "./test_model_{}.zip".format(request.node.name)
  
      try:
          env = DummyVecEnv([lambda: IdentityEnv(10)])
  
          # Create and save model with default MLP policy
          model = model_class(policy=MlpPolicy, env=env)
          model.save(model_fname)
  
          del model, env
  
          # Corrupt "policy" serialization in the file
          data_file = zipfile.ZipFile(model_fname, "r")
          # Load all data (can't just update one file in the archive)
          parameter_list = data_file.read("parameter_list")
          parameters = data_file.read("parameters")
          class_data = json.loads(data_file.read("data").decode())
          data_file.close()
  
          # Corrupt serialization of the "policy"
          class_data["policy"][":serialized:"] = (
              "Adding this should break serialization"
              + class_data["policy"][":serialized:"]
          )
  
          # And dump everything back to the model file
          data_file = zipfile.ZipFile(model_fname, "w")
          data_file.writestr("data", json.dumps(class_data))
          data_file.writestr("parameter_list", parameter_list)
          data_file.writestr("parameters", parameters)
          data_file.close()
  
          # Try loading the model. This should
          # result in an error
          with pytest.raises(RuntimeError):
              model = model_class.load(model_fname)
  
          # Load model with custom objects ("custom" MlpPolicy)
          # and it should work fine.
          # Note: We could load model with just vanilla
          #       MlpPolicy, too.
          model = model_class.load(
              model_fname, custom_objects={"policy": CustomMlpPolicy}
          )
  
          # Make sure we loaded custom MLP policy
          assert model.policy == CustomMlpPolicy
          del model
  
      finally:
          if os.path.exists(model_fname):
              os.remove(model_fname)
  
file_path: stable-baselines/tests/test_save.py
imports: import os
  from io import BytesIO
  import json
  import zipfile
  import pytest
  import numpy
  from stable_baselines import A2C, ACER, ACKTR, DQN, PPO1, PPO2, TRPO
  from stable_baselines.common.identity_env import IdentityEnv
  from stable_baselines.common.vec_env import DummyVecEnv
  from stable_baselines.common.evaluation import evaluate_policy
  from stable_baselines.common.policies import MlpPolicy, FeedForwardPolicy
label_desc: Use image node_id 2 for calling a global function with example usage: test_save_custom_objects(request, model_class) without return types