energy and force added as attributes to data objects for qm7x dataset (…

…ORNL#259) * energy and force added as attributes to data objects for qm7x dataset * energy explicitly added as attribute to data objects --------- Co-authored-by: Massimiliano Lupo Pasini <[email protected]>
allaffa · Jun 18, 2024 · ecd3135 · ecd3135
1 parent 80f0691
commit ecd3135
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/examples/qm7x/train.py b/examples/qm7x/train.py
@@ -194,17 +194,17 @@ def hdf5_to_graph(self, fMOL, molid):
                     forces
                 ), f"qm7x dataset - molid:{molid} - confid:{confid} - L2-norm of atomic forces exceeds {self.forces_norm_threshold}"
 
+                if self.energy_per_atom:
+                    energy = EPBE0 / natoms
+                else:
+                    energy = EPBE0
+
                 # data = Data(
                 #    pos=xyz, x=Z, molid=molid, confid=confid
                 # )
-                data = Data(pos=xyz, x=Z)
+                data = Data(pos=xyz, x=Z, force=forces, energy=energy, y=energy)
                 data.x = torch.cat((data.x, xyz, forces, hCHG, hVDIP, hRAT), dim=1)
 
-                if self.energy_per_atom:
-                    data.y = EPBE0 / natoms
-                else:
-                    data.y = EPBE0
-
                 data = create_graph_fromXYZ(data)
 
                 # Add edge length as edge feature