diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index ad5be7c709b0f090a3ff1d0705f73b49c283e848..a10c58e6f5eac8ebd04990505894a917c2212b3f 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -20,6 +20,8 @@ class TreeObsForRailEnv(ObservationBuilder):
     network to simplify the representation of the state of the environment for each agent.
     """
 
+    observation_dim = 9
+
     def __init__(self, max_depth, predictor=None):
         self.max_depth = max_depth
 
@@ -29,7 +31,6 @@ class TreeObsForRailEnv(ObservationBuilder):
         for i in range(self.max_depth + 1):
             size += pow4
             pow4 *= 4
-        self.observation_dim = 9
         self.observation_space = [size * self.observation_dim]
         self.location_has_agent = {}
         self.location_has_agent_direction = {}
diff --git a/tests/test_flatland_envs_predictions.py b/tests/test_flatland_envs_predictions.py
index 96a4b9589093d8635006dac8bea8e902d449600a..c90f91a041b16cee2dc55a58562d34a0b9100560 100644
--- a/tests/test_flatland_envs_predictions.py
+++ b/tests/test_flatland_envs_predictions.py
@@ -265,8 +265,8 @@ def test_shortest_path_predictor_conflicts(rendering=False):
     pp.pprint(tree_0)
 
     # check the expectations
-    expected_conflicts_0 = [('F','R')]
-    expected_conflicts_1 = [('F','L')]
+    expected_conflicts_0 = [('F', 'R')]
+    expected_conflicts_1 = [('F', 'L')]
     _check_expected_conflicts(expected_conflicts_0, obs_builder, tree_0, "agent[0]: ")
     _check_expected_conflicts(expected_conflicts_1, obs_builder, tree_1, "agent[1]: ")
 
diff --git a/tox.ini b/tox.ini
index 2047166aace141680561623c7edf4d11a2b4534f..644f5253e103c99367096eff0d44d1456ae88ba7 100644
--- a/tox.ini
+++ b/tox.ini
@@ -113,6 +113,8 @@ deps =
     -r{toxinidir}/requirements_dev.txt
     -r{toxinidir}/requirements_continuous_integration.txt
 commands =
+; install current version of flatland to be used by notebooks
+    sh -c 'python setup.py'
 ; run tests from subfolder to ensure that resources are accessed via resources and not via relative paths
     sh -c 'mkdir -p {envtmpdir}/6f59bc68108c3895b1828abdd04b9a06'
     sh -c 'jupyter nbextension install --py --sys-prefix widgetsnbextension'