diff --git a/maze_transformer/training/train_model.py b/maze_transformer/training/train_model.py
index b98b66ad..ce18d9dc 100644
--- a/maze_transformer/training/train_model.py
+++ b/maze_transformer/training/train_model.py
@@ -194,7 +194,7 @@ def train_model(
         output_dir=output_path,
         device=device,
         val_dataset=val_dataset,
-    )    
+    )
 
     return TrainingResult(
         output_path=output_path,