Relax config (#758)

* update example config per PR714 * no need to support this * y->energy
FAIR-Chem · Jul 12, 2024 · 1442177 · 1442177
1 parent 2e8578a
commit 1442177
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 92 deletions.
diff --git a/configs/ocp_example.yml b/configs/ocp_example.yml
@@ -71,6 +71,13 @@ dataset:
   test:
     # Directory containing test set LMDBs
     src: data/s2ef/all/test_id/
+  relax:
+    # Path to initial structures to run relaxations on. Same as the IS2RE set.
+    src: data/is2re/all/test_id/data.lmdb
+    # To shard a dataset into smaller subsets, define the total_shards desired
+    # and the shard a particular process to see.
+    total_shards: 1                                                             # int (optional)
+    shard: 0                                                                    # int (optional)
 
 task:
   # This is an argument used for checkpoint loading. By default it is True and loads
@@ -92,13 +99,6 @@ task:
   relaxation_fmax: 0.02
   # Whether to save out the positions.
   write_pos: True                                                               # True or False
-  # Path to initial structures to run relaxations on. Same as the IS2RE set.
-  relax_dataset:
-    src: data/is2re/all/test_id/data.lmdb
-    # To shard a dataset into smaller subsets, define the total_shards desired
-    # and the shard a particular process to see.
-    total_shards: 1                                                             # int (optional)
-    shard: 0                                                                    # int (optional)
   relax_opt:
     name: lbfgs
     maxstep: 0.2

diff --git a/configs/s2ef/2M/dimenet_plus_plus/dpp_relax.yml b/configs/s2ef/2M/dimenet_plus_plus/dpp_relax.yml
diff --git a/src/fairchem/core/common/relaxation/ase_utils.py b/src/fairchem/core/common/relaxation/ase_utils.py
@@ -43,7 +43,7 @@ def batch_to_atoms(batch):
     positions = torch.split(batch.pos, natoms)
     tags = torch.split(batch.tags, natoms)
     cells = batch.cell
-    energies = batch.y.view(-1).tolist()
+    energies = batch.energy.view(-1).tolist()
 
     atoms_objects = []
     for idx in range(n_systems):