ds bug fix

Author: Tetracarbonylnickel

apax/train/run.py

@@ -85,14 +85,6 @@ def initialize_dataset(config, raw_ds, calc_stats: bool = True):
         energy_unit=config.data.energy_unit,
     )
 
-    dataset = TFPipeline(
-        inputs,
-        labels,
-        config.n_epochs,
-        config.data.batch_size,
-        buffer_size=config.data.shuffle_buffer_size,
-    )
-
     if calc_stats:
         ds_stats = compute_scale_shift_parameters(
             inputs,
@@ -103,6 +95,15 @@ def initialize_dataset(config, raw_ds, calc_stats: bool = True):
             config.data.scale_options,
         )
 
+    dataset = TFPipeline(
+        inputs,
+        labels,
+        config.n_epochs,
+        config.data.batch_size,
+        buffer_size=config.data.shuffle_buffer_size,
+    )
+
+    if calc_stats:
         return dataset, ds_stats
     else:
         return dataset

apax/utils/convert.py

@@ -109,7 +109,6 @@ def atoms_to_arrays(
 
         inputs["ragged"]["numbers"].append(atoms.numbers)
         inputs["fixed"]["n_atoms"].append(len(atoms))
-
         for key, val in atoms.calc.results.items():
             if key == "forces":
                 labels["ragged"][key].append(

tests/unit_tests/data/test_statistics.py

@@ -24,7 +24,7 @@ def test_energy_per_element():
         atoms.calc = SinglePointCalculator(atoms, energy=energy)
 
     labels = {
-        "ragged": {
+        "fixed": {
             "energy": [atoms.get_potential_energy() for atoms in atoms_list],
         }
     }