diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index 4fad2c01c4d353b6ba458a94d6fb10085213eb06..bde43bf8f6f998b0cf3f0c91925b1ba6b11dd622 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -396,6 +396,9 @@ class TreeObsForRailEnv(ObservationBuilder):
                 observation = observation + branch_observation
                 if len(branch_visited) != 0:
                     visited.union(branch_visited)
+                    print(visited)
+                    print(branch_visited)
+                    print("-------")
             else:
                 num_cells_to_fill_in = 0
                 pow4 = 1
diff --git a/tests/test_environments.py b/tests/test_environments.py
index 4c55eac7afb44d95f0e49d665eeeb4bc36becea9..9c7b53b9b5876a99d7deea20da10816d81f02b65 100644
--- a/tests/test_environments.py
+++ b/tests/test_environments.py
@@ -7,7 +7,7 @@ from flatland.envs.generators import rail_from_GridTransitionMap_generator
 from flatland.envs.generators import complex_rail_generator
 from flatland.core.transitions import Grid4Transitions
 from flatland.core.transition_map import GridTransitionMap
-from flatland.core.env_observation_builder import GlobalObsForRailEnv
+from flatland.envs.observations import GlobalObsForRailEnv
 from flatland.envs.agent_utils import EnvAgent
 
 """Tests for `flatland` package."""