Allow no energy data for MACEField.

Author: Neutrino155

mace/cli/run_train.py

@@ -309,6 +309,7 @@ def run(args) -> None:
                     args.pseudolabel_replay
                     and args.multiheads_finetuning
                     and head_config.head_name == "pt_head"
+                    or args.model == "MACEField"
                 ),
             )
             head_config.collections = SubsetCollection(

mace/tools/scripts_utils.py

@@ -515,6 +515,15 @@ def get_atomic_energies(E0s, train_collection, z_table) -> dict:
                 atomic_energies_dict = data.compute_average_E0s(
                     train_collection, z_table
                 )
+                atomic_energies_dict = {
+                    key:value 
+                    if 
+                        not np.isnan(value) 
+                    else 
+                        0.0 
+                    for 
+                        key, value in atomic_energies_dict.items()
+                }
             except Exception as e:
                 raise RuntimeError(
                     f"Could not compute average E0s if no training xyz given, error {e} occured"