# Buggy code
def load(
    cls,
    path: Union[str, pathlib.Path, io.BufferedIOBase],
    env: Optional[GymEnv] = None,
    device: Union[th.device, str] = "auto",
    **kwargs,
) -> "BaseAlgorithm":
    """
    Load the model from a zip-file

    :param path: path to the file (or a file-like) where to
        load the agent from
    :param env: the new environment to run the loaded model on
        (can be None if you only need prediction from a trained model) has priority over any saved environment
    :param device: Device on which the code should run.
    :param kwargs: extra arguments to change the model when loading
    """
    data, params, pytorch_variables = load_from_zip_file(path, device=device)

    # Remove stored device information and replace with ours
    if "policy_kwargs" in data:
        if "device" in data["policy_kwargs"]:
            del data["policy_kwargs"]["device"]

    if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]:
        raise ValueError(
            f"The specified policy kwargs do not equal the stored policy kwargs."
            f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}"
        )

    if "observation_space" not in data or "action_space" not in data:
        raise KeyError(
            "The observation_space and action_space were not given, can't verify new environments"
        )

    if env is not None:
        # Wrap first if needed
        cls._wrap_env(env, data["verbose"])
        # Check if given env is valid
        check_for_correct_spaces(env, data["observation_space"], data["action_space"])
    else:
        # Use stored env, if one exists. If not, continue as is (can be used for predict)
        if "env" in data:
            env = data["env"]

    # noinspection PyArgumentList
    model = cls(
        policy=data["policy_class"],
        env=env,
        device=device,
        _init_setup_model=False,  # pytype: disable=not-instantiable,wrong-keyword-args
    )

    # load parameters
    model.__dict__.update(data)
    model.__dict__.update(kwargs)
    model._setup_model()

    # put state_dicts back in place
    model.set_parameters(params, exact_match=True, device=device)

    # put other pytorch variables back in place
    if pytorch_variables is not None:
        for name in pytorch_variables:
            recursive_setattr(model, name, pytorch_variables[name])

    # Sample gSDE exploration matrix, so it uses the right device
    # see issue #44
    if model.use_sde:
        model.policy.reset_noise()  # pytype: disable=attribute-error
    return model


# Traceback (ValueError)
Box(36, 36, 4)
Box(4, 36, 36)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-176-a54409b23a82> in <module>
1 model.save("custom_env")
----> 2 model = PPO.load("custom_env", custom_env)

/opt/conda/lib/python3.7/site-packages/stable_baselines3/common/base_class.py in load(cls, path, env, device, **kwargs)
595             cls._wrap_env(env, data["verbose"])
596             # Check if given env is valid
--> 597             check_for_correct_spaces(env, data["observation_space"], data["action_space"])
598         else:
599             # Use stored env, if one exists. If not, continue as is (can be used for predict)

/opt/conda/lib/python3.7/site-packages/stable_baselines3/common/utils.py in check_for_correct_spaces(env, observation_space, action_space)
204     """
205     if observation_space != env.observation_space:
--> 206         raise ValueError(f"Observation spaces do not match: {observation_space} != {env.observation_space}")
207     if action_space != env.action_space:
208         raise ValueError(f"Action spaces do not match: {action_space} != {env.action_space}")

ValueError: Observation spaces do not match: Box(4, 36, 36) != Box(36, 36, 4)

# Fixed code
def load(
    cls,
    path: Union[str, pathlib.Path, io.BufferedIOBase],
    env: Optional[GymEnv] = None,
    device: Union[th.device, str] = "auto",
    **kwargs,
) -> "BaseAlgorithm":
    """
    Load the model from a zip-file

    :param path: path to the file (or a file-like) where to
        load the agent from
    :param env: the new environment to run the loaded model on
        (can be None if you only need prediction from a trained model) has priority over any saved environment
    :param device: Device on which the code should run.
    :param kwargs: extra arguments to change the model when loading
    """
    data, params, pytorch_variables = load_from_zip_file(path, device=device)

    # Remove stored device information and replace with ours
    if "policy_kwargs" in data:
        if "device" in data["policy_kwargs"]:
            del data["policy_kwargs"]["device"]

    if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]:
        raise ValueError(
            f"The specified policy kwargs do not equal the stored policy kwargs."
            f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}"
        )

    if "observation_space" not in data or "action_space" not in data:
        raise KeyError(
            "The observation_space and action_space were not given, can't verify new environments"
        )

    if env is not None:
        # Wrap first if needed
        env = cls._wrap_env(env, data["verbose"])
        # Check if given env is valid
        check_for_correct_spaces(env, data["observation_space"], data["action_space"])
    else:
        # Use stored env, if one exists. If not, continue as is (can be used for predict)
        if "env" in data:
            env = data["env"]

    # noinspection PyArgumentList
    model = cls(
        policy=data["policy_class"],
        env=env,
        device=device,
        _init_setup_model=False,  # pytype: disable=not-instantiable,wrong-keyword-args
    )

    # load parameters
    model.__dict__.update(data)
    model.__dict__.update(kwargs)
    model._setup_model()

    # put state_dicts back in place
    model.set_parameters(params, exact_match=True, device=device)

    # put other pytorch variables back in place
    if pytorch_variables is not None:
        for name in pytorch_variables:
            recursive_setattr(model, name, pytorch_variables[name])

    # Sample gSDE exploration matrix, so it uses the right device
    # see issue #44
    if model.use_sde:
        model.policy.reset_noise()  # pytype: disable=attribute-error
    return model