diff --git a/mace/cli/run_train.py b/mace/cli/run_train.py
index 9b484d7f..7f1a5e74 100644
--- a/mace/cli/run_train.py
+++ b/mace/cli/run_train.py
@@ -364,7 +364,7 @@ def run(args: argparse.Namespace) -> None:
                         z_table_foundation.z_to_index(z)
                     ].item()
                     for z in z_table.zs
-                }         
+                }
             else:
                 atomic_energies_dict[head_config.head_name] = get_atomic_energies(head_config.E0s, None, head_config.z_table)
         else: