diff --git a/examples/custom_railmap_example.py b/examples/custom_railmap_example.py
index 6dc157f33ab90d5efa06dc45c914a8de79def7e8..a76b9c974fd04dba84622217f347a591498b818a 100644
--- a/examples/custom_railmap_example.py
+++ b/examples/custom_railmap_example.py
@@ -2,8 +2,8 @@ import random
 
 import numpy as np
 
-from flatland.core.transition_map import GridTransitionMap
 from flatland.core.grid.rail_env_grid import RailEnvTransitions
+from flatland.core.transition_map import GridTransitionMap
 from flatland.envs.rail_env import RailEnv
 from flatland.utils.rendertools import RenderTool
 
@@ -17,7 +17,7 @@ def custom_rail_generator():
         grid_map = GridTransitionMap(width=width, height=height, transitions=rail_trans)
         rail_array = grid_map.grid
         rail_array.fill(0)
-        new_tran = rail_trans.set_transition(1,1,1,1)
+        new_tran = rail_trans.set_transition(1, 1, 1, 1)
         print(new_tran)
         agents_positions = []
         agents_direction = []
diff --git a/tests/simple_rail.py b/tests/simple_rail.py
index fb739f29942d0552b41b602046030c76031e7539..28978ca3e5b958a51c578f6be5e8c87b77baaa97 100644
--- a/tests/simple_rail.py
+++ b/tests/simple_rail.py
@@ -1,11 +1,12 @@
-import numpy as np
 from typing import Tuple
 
+import numpy as np
+
 from flatland.core.grid.grid4 import Grid4Transitions
 from flatland.core.transition_map import GridTransitionMap
 
 
-def make_simple_rail() -> Tuple[GridTransitionMap,np.array]:
+def make_simple_rail() -> Tuple[GridTransitionMap, np.array]:
     # We instantiate a very simple rail network on a 7x10 grid:
     #        |
     #        |