diff --git a/torch_training/bla.py b/torch_training/bla.py
index 80ec308c2b6bc5498d9198e2a03e562b02e7c96d..ec81697ac1a41b6ca2ec12803b6c6817578ea6ed 100644
--- a/torch_training/bla.py
+++ b/torch_training/bla.py
@@ -110,13 +110,13 @@ def main(argv):
             max_steps = int(3 * (env.height + env.width))
             agent_obs = [None] * env.get_num_agents()
             agent_next_obs = [None] * env.get_num_agents()
-        # # Reset environment
-        # if file_load:
-        #     obs = env.reset(False, False)
-        # else:
-        #     obs = env.reset(True, True)
-        # if demo:
-        #     env_renderer.set_new_rail()
+        # Reset environment
+        if file_load:
+            obs = env.reset(False, False)
+        else:
+            obs = env.reset(True, True)
+        if demo:
+            env_renderer.set_new_rail()
         # obs_original = obs.copy()
         # final_obs = obs.copy()
         # final_obs_next = obs.copy()