diff --git a/tests/test_save_load.py b/tests/test_save_load.py
index 6368127..e184c6e 100644
--- a/tests/test_save_load.py
+++ b/tests/test_save_load.py
@@ -10,9 +10,9 @@ from torchy_baselines.common.identity_env import IdentityEnvBox
 
 MODEL_LIST = [
     PPO,
-    A2C,
-    TD3,
-    SAC,
+    #A2C,
+    #TD3,
+    #SAC,
 ]
 
 
@@ -30,6 +30,7 @@ def test_save_load(model_class):
 
     # create model
     model = model_class('MlpPolicy', env, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
+    model.learn(total_timesteps=1000, eval_freq=500)
 
     # Get dictionary of current parameters
     params = deepcopy(model.get_policy_parameters())
@@ -50,7 +51,6 @@ def test_save_load(model_class):
     params = new_params
 
     # Check
-    model.learn(total_timesteps=1000, eval_freq=500)
     model.save("test_save.zip")
     del model
     model = model_class.load("test_save")
@@ -66,7 +66,14 @@ def test_save_load(model_class):
     new_opt_params = model.get_opt_parameters()
     # check if keys are the same
     assert opt_params.keys() == new_opt_params.keys()
-    # check if values are the same: don't know how to to that
+    # check if values are the same: only tested for Adam and RMSProp so far
+    for optimizer,opt_state in opt_params.items():
+        for step_entry, entry_dict in opt_state['state'].items():
+            for value_key,value in entry_dict.items():
+                print(value == new_opt_params[optimizer][step_entry][value_key])
+
+
+
 
     # check if learn still works
     model.learn(total_timesteps=1000, eval_freq=500)
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 3384775..597d7fe 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -503,7 +503,7 @@ class BaseRLModel(object):
                     with file_.open(file_name + '.pth', mode="w") as opt_param_file:
                         th.save(dict, opt_param_file)
 
-    def save(self, path):
+    def save(self, path, include=None):#TODO
         """
         saves all the params from init and pytorch params in a file for continuous learning
 
@@ -514,16 +514,4 @@ class BaseRLModel(object):
         data.pop("replay_buffer")
         params_to_save = self.get_policy_parameters()
         opt_params_to_save = self.get_opt_parameters()
-        self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
-
-    def save_with_replay_buffer(self, path):
-        """
-        saves all the params from init and pytorch params in a file for continuous learning
-
-        :param path: path to the file where the data should be saved
-        :return:
-        """
-        data = self.__dict__
-        params_to_save = self.get_policy_parameters()
-        opt_params_to_save = self.get_opt_parameters()
-        self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
+        self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
\ No newline at end of file
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index c4e0838..7582151 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -124,14 +124,14 @@ class PPO(BaseRLModel):
         if self.clip_range_vf is not None:
             self.clip_range_vf = get_schedule_fn(self.clip_range_vf)
 
-    def select_action(self, observation):
+    def select_action(self, observation,deterministic=False):
         # Normally not needed
         observation = np.array(observation)
         with th.no_grad():
             observation = th.FloatTensor(observation.reshape(1, -1)).to(self.device)
-            return self.policy.actor_forward(observation, deterministic=False)
+            return self.policy.actor_forward(observation, deterministic)
 
-    def predict(self, observation, state=None, mask=None, deterministic=True):
+    def predict(self, observation, state=None, mask=None, deterministic=False):
         """
         Get the model's action from an observation