diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 7260f11..8eab899 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -404,7 +404,7 @@ def load_from_zip_file( # Load the parameters with the right ``map_location``. # Remove ".pth" ending with splitext th_object = th.load(file_content, map_location=device) - if file_path == "pytorch_variables.pth": + if file_path == "pytorch_variables.pth" or file_path == "tensors.pth": # PyTorch variables (not state_dicts) pytorch_variables = th_object else: diff --git a/stable_baselines3/tqc/tqc.py b/stable_baselines3/tqc/tqc.py index 34f370b..750c97b 100644 --- a/stable_baselines3/tqc/tqc.py +++ b/stable_baselines3/tqc/tqc.py @@ -468,7 +468,7 @@ class TQC(OffPolicyAlgorithm): reset_num_timesteps=reset_num_timesteps, ) - def excluded_save_params(self) -> List[str]: + def _excluded_save_params(self) -> List[str]: """ Returns the names of the parameters that should be excluded by default when saving the model. @@ -478,14 +478,14 @@ class TQC(OffPolicyAlgorithm): # Exclude aliases return super(TQC, self).excluded_save_params() + ["actor", "critic", "critic_target"] - def get_torch_variables(self) -> Tuple[List[str], List[str]]: + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: """ cf base class """ state_dicts = ["policy", "actor.optimizer", "critic.optimizer"] - saved_tensors = ["log_ent_coef"] + saved_pytorch_variables = ["log_ent_coef"] if self.ent_coef_optimizer is not None: state_dicts.append("ent_coef_optimizer") else: - saved_tensors.append("ent_coef_tensor") - return state_dicts, saved_tensors + saved_pytorch_variables.append("ent_coef_tensor") + return state_dicts, saved_pytorch_variables