diff --git a/tests/test_save_load.py b/tests/test_save_load.py
index bae9e0e..eef023d 100644
--- a/tests/test_save_load.py
+++ b/tests/test_save_load.py
@@ -89,8 +89,8 @@ def test_save_load(model_class):
 
     # check if model still selects the same actions
     new_selected_actions = [model.predict(observation, deterministic=True) for observation in observations]
-    for i in range(len(selected_actions)):
-        assert selected_actions[i] == new_selected_actions[i]
+    # for i in range(len(selected_actionsselected_actions)):
+    assert np.allclose(selected_actions, new_selected_actions)
 
     # check if learn still works
     model.learn(total_timesteps=1000, eval_freq=500)
@@ -99,6 +99,11 @@ def test_save_load(model_class):
     os.remove("test_save.zip")
 
 
+@pytest.mark.parametrize("model_class", MODEL_LIST)
+def test_set_env(model_class):
+    pass
+
+
 @pytest.mark.parametrize("model_class", MODEL_LIST)
 def test_exclude_include_saved_params(model_class):
     """
@@ -126,6 +131,5 @@ def test_exclude_include_saved_params(model_class):
     model = model_class.load("test_save")
     assert model.verbose == 2
 
-
     # clear file from os
     os.remove("test_save.zip")
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index 384ad27..d10dec6 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -166,6 +166,22 @@ class BaseRLModel(object):
         """
         return self.env
 
+    @staticmethod
+    def check_env(env, observation_space, action_space):
+        """
+        Checks the validity of the environment and returns if it is coherent
+        Checked parameters:
+         - observation_space
+         - action_space
+        :return: (bool) True if environment seems to be coherent
+        """
+        if observation_space != env.observation_space:
+            return False
+        if action_space != env.action_space:
+            return False
+        # return true if no check failed
+        return True
+
     def set_env(self, env):
         """
         Checks the validity of the environment, and if it is coherent, set it as the current environment.
@@ -175,21 +191,20 @@ class BaseRLModel(object):
 
         :param env: (Gym Environment) The environment for learning a policy
         """
-
-        if self.observation_space != env.observation_space:
-            raise ValueError("The given environment has a observation_space that doesn't fit the current model")
-
-        if self.action_space != env.action_space:
-            raise ValueError("The given environment has a action_space that doesn't fit the current model")
+        if self.check_env(env, self.observation_space, self.action_space) is False:
+            raise ValueError("Given environment is not compatible with model")
         # if all fits save new env
         self.env = env
+        # and update observation and action space
+        self.observation_space = env.observation_space
+        self.action_space = env.action_space
 
     def get_parameter_list(self):
         """
         Returns policy and optimizer parameters as a tuple
         :return: (dict,dict) policy_parameters, opt_parameters
         """
-        return self.get_policy_parameters(),self.get_opt_parameters()
+        return self.get_policy_parameters(), self.get_opt_parameters()
 
     def get_policy_parameters(self):
         """
@@ -286,15 +301,22 @@ class BaseRLModel(object):
                              "Stored kwargs: {}, specified kwargs: {}".format(data['policy_kwargs'],
                                                                               kwargs['policy_kwargs']))
 
+        # check if observation space and action space is given
+        if ("observation_space" not in data or "action_space" not in data) and "env" not in data:
+            raise ValueError("The observation_space and action_space was not given, can't verify new environments")
+        # check if given env is valid
+        if env is not None and cls.check_env(env, data["observation_space"], data["action_space"]) is False:
+            raise ValueError("The given environment does not comply to the model")
+        # if no new env was given use stored env if possible
         if env is None and "env" in data:
             env = data["env"]
-        if env is not None:
-            model = cls(policy=data["policy_class"], env=env, _init_setup_model=True)
-        else:
-            model = cls(policy=data["policy_class"], env=env, _init_setup_model=False)
+
+        # first create model, but only setup if a env was given
+        model = cls(policy=data["policy_class"], env=env, _init_setup_model=env is not None)
+
+        # load parameters
         model.__dict__.update(data)
         model.__dict__.update(kwargs)
-        model.set_env(env)
         model.load_parameters(params, opt_params)
         return model
 
@@ -342,7 +364,7 @@ class BaseRLModel(object):
 
                 # check for all other .pth files
                 other_file = [file_name for file_name in namelist if
-                               os.path.splitext(file_name)[1] == ".pth" and file_name != "params.pth"]
+                              os.path.splitext(file_name)[1] == ".pth" and file_name != "params.pth"]
                 # if there are any other files which end with .pth and aren't "params.pth"
                 # assume that they each are optimizer parameters
                 if len(other_file) > 0: