change order of key mapping transform (#993)

* change order of key mapping

* move transform to before renaming

* fix lint

Author: misko

src/fairchem/core/datasets/ase_datasets.py

@@ -135,12 +135,12 @@ def __getitem__(self, idx):
         if self.a2g.r_energy is True and self.lin_ref is not None:
             data_object.energy -= sum(self.lin_ref[data_object.atomic_numbers.long()])
 
-        if self.key_mapping is not None:
-            data_object = rename_data_object_keys(data_object, self.key_mapping)
-
         # Transform data object
         data_object = self.transforms(data_object)
 
+        if self.key_mapping is not None:
+            data_object = rename_data_object_keys(data_object, self.key_mapping)
+
         if self.config.get("include_relaxed_energy", False):
             data_object.energy_relaxed = self.get_relaxed_energy(self.ids[idx])
 

src/fairchem/core/datasets/lmdb_dataset.py

@@ -145,10 +145,12 @@ def __getitem__(self, idx: int) -> T_co:
             )
             data_object = pyg2_data_transform(pickle.loads(datapoint_pickled))
 
+        data_object = self.transforms(data_object)
+
         if self.key_mapping is not None:
             data_object = rename_data_object_keys(data_object, self.key_mapping)
 
-        return self.transforms(data_object)
+        return data_object
 
     def connect_db(self, lmdb_path: Path | None = None) -> lmdb.Environment:
         return lmdb.open(

src/fairchem/core/datasets/oc22_lmdb_dataset.py

@@ -192,9 +192,6 @@ def __getitem__(self, idx):
             lin_energy = sum(self.lin_ref[data_object.atomic_numbers.long()])
             data_object[attr] -= lin_energy
 
-        if self.key_mapping is not None:
-            data_object = rename_data_object_keys(data_object, self.key_mapping)
-
         # to jointly train on oc22+oc20, need to delete these oc20-only attributes
         # ensure otf_graph=1 in your model configuration
         if "edge_index" in data_object:
@@ -204,7 +201,12 @@ def __getitem__(self, idx):
         if "distances" in data_object:
             del data_object.distances
 
-        return self.transforms(data_object)
+        data_object = self.transforms(data_object)
+
+        if self.key_mapping is not None:
+            data_object = rename_data_object_keys(data_object, self.key_mapping)
+
+        return data_object
 
     def connect_db(self, lmdb_path=None):
         return lmdb.open(