diff --git a/configs/s2ef/2M/mace/012702_128x0e_128x1o_128x2e.yml b/configs/s2ef/2M/mace/012702_128x0e_128x1o_128x2e.yml
new file mode 100644
index 0000000000..4f66a00b8c
--- /dev/null
+++ b/configs/s2ef/2M/mace/012702_128x0e_128x1o_128x2e.yml
@@ -0,0 +1,78 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  avg_num_neighbors: 35
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 3
+  num_elements: 83
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 128x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  atomic_inter_shift: -0.75544
+  atomic_inter_scale: 2.88732
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # Default in MACE (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  # was 1.e-2.
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/013101_correct_scaleshift_weight_decay.yml b/configs/s2ef/2M/mace/013101_correct_scaleshift_weight_decay.yml
new file mode 100644
index 0000000000..0cd8876c18
--- /dev/null
+++ b/configs/s2ef/2M/mace/013101_correct_scaleshift_weight_decay.yml
@@ -0,0 +1,80 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  avg_num_neighbors: 35
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 3
+  num_elements: 83
+  hidden_irreps: 128x0e + 128x1o
+  MLP_irreps: 128x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.0106
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4263
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 16
+  eval_batch_size: 16
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # Default in MACE (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  # was 1.e-2.
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 1.e-10
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/013102_128x0e_128x1o_128x2e.yml b/configs/s2ef/2M/mace/013102_128x0e_128x1o_128x2e.yml
new file mode 100644
index 0000000000..bc3c87dce3
--- /dev/null
+++ b/configs/s2ef/2M/mace/013102_128x0e_128x1o_128x2e.yml
@@ -0,0 +1,80 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  avg_num_neighbors: 35
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 3
+  num_elements: 83
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 128x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.0106
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4263
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # Default in MACE (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  # was 1.e-2.
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 1.e-10
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/013103_h256.yml b/configs/s2ef/2M/mace/013103_h256.yml
new file mode 100644
index 0000000000..cc159e7482
--- /dev/null
+++ b/configs/s2ef/2M/mace/013103_h256.yml
@@ -0,0 +1,80 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  avg_num_neighbors: 35
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 3
+  num_elements: 83
+  hidden_irreps: 256x0e + 256x1o
+  MLP_irreps: 256x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.0106
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4263
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # Default in MACE (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  # was 1.e-2.
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 1.e-10
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/032202_bugfixed.yml b/configs/s2ef/2M/mace/032202_bugfixed.yml
new file mode 100644
index 0000000000..8a1064b96e
--- /dev/null
+++ b/configs/s2ef/2M/mace/032202_bugfixed.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 8.0
+  max_neighbors: 500
+  avg_num_neighbors: 72.44458770751953
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.033861486108663866
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/032501_wd1e-10.yml b/configs/s2ef/2M/mace/032501_wd1e-10.yml
new file mode 100644
index 0000000000..a0c0bfae05
--- /dev/null
+++ b/configs/s2ef/2M/mace/032501_wd1e-10.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 8.0
+  max_neighbors: 500
+  avg_num_neighbors: 72.44458770751953
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.033861486108663866
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 1.e-10
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/032502_ef_lr_switch_ft.yml b/configs/s2ef/2M/mace/032502_ef_lr_switch_ft.yml
new file mode 100644
index 0000000000..09402d7298
--- /dev/null
+++ b/configs/s2ef/2M/mace/032502_ef_lr_switch_ft.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 8.0
+  max_neighbors: 500
+  avg_num_neighbors: 72.44458770751953
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.033861486108663866
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1
+  energy_coefficient: 1000
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/032503_MLP_irreps_128x0e.yml b/configs/s2ef/2M/mace/032503_MLP_irreps_128x0e.yml
new file mode 100644
index 0000000000..7ad6d999ff
--- /dev/null
+++ b/configs/s2ef/2M/mace/032503_MLP_irreps_128x0e.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 8.0
+  max_neighbors: 500
+  avg_num_neighbors: 72.44458770751953
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.033861486108663866
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/033001_bugfixed_nbrs30.yml b/configs/s2ef/2M/mace/033001_bugfixed_nbrs30.yml
new file mode 100644
index 0000000000..6874afb641
--- /dev/null
+++ b/configs/s2ef/2M/mace/033001_bugfixed_nbrs30.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/033002_bugfixed_nbrs20.yml b/configs/s2ef/2M/mace/033002_bugfixed_nbrs20.yml
new file mode 100644
index 0000000000..4ea77f6d74
--- /dev/null
+++ b/configs/s2ef/2M/mace/033002_bugfixed_nbrs20.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 20
+  avg_num_neighbors: 20
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/033003_bugfixed_nbrs30_allatoms.yml b/configs/s2ef/2M/mace/033003_bugfixed_nbrs30_allatoms.yml
new file mode 100644
index 0000000000..a5f75228b2
--- /dev/null
+++ b/configs/s2ef/2M/mace/033003_bugfixed_nbrs30_allatoms.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: False
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/033101_nbrs30_maxell4.yml b/configs/s2ef/2M/mace/033101_nbrs30_maxell4.yml
new file mode 100644
index 0000000000..9428cf1c56
--- /dev/null
+++ b/configs/s2ef/2M/mace/033101_nbrs30_maxell4.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 4
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/033102_nbrs30_adam_wd5e-10.yml b/configs/s2ef/2M/mace/033102_nbrs30_adam_wd5e-10.yml
new file mode 100644
index 0000000000..f3caaaa7d3
--- /dev/null
+++ b/configs/s2ef/2M/mace/033102_nbrs30_adam_wd5e-10.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: Adam
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-10
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/033103_nbrs30_adam_wd5e-7.yml b/configs/s2ef/2M/mace/033103_nbrs30_adam_wd5e-7.yml
new file mode 100644
index 0000000000..b560634f02
--- /dev/null
+++ b/configs/s2ef/2M/mace/033103_nbrs30_adam_wd5e-7.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: Adam
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040302_nbrs30_gaussian.yml b/configs/s2ef/2M/mace/040302_nbrs30_gaussian.yml
new file mode 100644
index 0000000000..a7ea5bbd44
--- /dev/null
+++ b/configs/s2ef/2M/mace/040302_nbrs30_gaussian.yml
@@ -0,0 +1,91 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040501_nbrs30_adam_wd0.yml b/configs/s2ef/2M/mace/040501_nbrs30_adam_wd0.yml
new file mode 100644
index 0000000000..7a7b963640
--- /dev/null
+++ b/configs/s2ef/2M/mace/040501_nbrs30_adam_wd0.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 6
+  eval_batch_size: 6
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: Adam
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 0.
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040502_nbrs30_gaussian_rbfh256.yml b/configs/s2ef/2M/mace/040502_nbrs30_gaussian_rbfh256.yml
new file mode 100644
index 0000000000..21fc910a1b
--- /dev/null
+++ b/configs/s2ef/2M/mace/040502_nbrs30_gaussian_rbfh256.yml
@@ -0,0 +1,92 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 128x0e + 128x1o + 128x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 12
+  eval_batch_size: 12
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040503_nbrs30_gaussian_h512_mlp128_rbfh256.yml b/configs/s2ef/2M/mace/040503_nbrs30_gaussian_h512_mlp128_rbfh256.yml
new file mode 100644
index 0000000000..9b53bd2747
--- /dev/null
+++ b/configs/s2ef/2M/mace/040503_nbrs30_gaussian_h512_mlp128_rbfh256.yml
@@ -0,0 +1,92 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040601_h1024_mlp128_rbfh256.yml b/configs/s2ef/2M/mace/040601_h1024_mlp128_rbfh256.yml
new file mode 100644
index 0000000000..7d64b9b9cd
--- /dev/null
+++ b/configs/s2ef/2M/mace/040601_h1024_mlp128_rbfh256.yml
@@ -0,0 +1,92 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 1024x0e + 1024x1o + 1024x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 2
+  eval_batch_size: 2
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040602_h512_mlp128_rbfh256_direct.yml b/configs/s2ef/2M/mace/040602_h512_mlp128_rbfh256_direct.yml
new file mode 100644
index 0000000000..bda823728d
--- /dev/null
+++ b/configs/s2ef/2M/mace/040602_h512_mlp128_rbfh256_direct.yml
@@ -0,0 +1,93 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.03386107087135315
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.4365839660167694
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040701_h512_mlp128_rbfh256_direct_intenergies.yml b/configs/s2ef/2M/mace/040701_h512_mlp128_rbfh256_direct_intenergies.yml
new file mode 100644
index 0000000000..ec4c10d82b
--- /dev/null
+++ b/configs/s2ef/2M/mace/040701_h512_mlp128_rbfh256_direct_intenergies.yml
@@ -0,0 +1,84 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040801_intenergies_maxell4.yml b/configs/s2ef/2M/mace/040801_intenergies_maxell4.yml
new file mode 100644
index 0000000000..369c1f54ac
--- /dev/null
+++ b/configs/s2ef/2M/mace/040801_intenergies_maxell4.yml
@@ -0,0 +1,84 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 4
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040802_wd1e-2.yml b/configs/s2ef/2M/mace/040802_wd1e-2.yml
new file mode 100644
index 0000000000..3764ff0177
--- /dev/null
+++ b/configs/s2ef/2M/mace/040802_wd1e-2.yml
@@ -0,0 +1,84 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 0.01
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/040803_allatoms_0.1_1_1.yml b/configs/s2ef/2M/mace/040803_allatoms_0.1_1_1.yml
new file mode 100644
index 0000000000..278bc4c8d3
--- /dev/null
+++ b/configs/s2ef/2M/mace/040803_allatoms_0.1_1_1.yml
@@ -0,0 +1,88 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  tag_specific_weights:
+    - 0.1
+    - 1.0
+    - 1.0
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+  loss_force: l2mae
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/041201_speciesagnostic.yml b/configs/s2ef/2M/mace/041201_speciesagnostic.yml
new file mode 100644
index 0000000000..30781c8365
--- /dev/null
+++ b/configs/s2ef/2M/mace/041201_speciesagnostic.yml
@@ -0,0 +1,87 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticInteractionBlock
+  interaction_cls: SpeciesAgnosticInteractionBlock
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/041801_specagn_withskip.yml b/configs/s2ef/2M/mace/041801_specagn_withskip.yml
new file mode 100644
index 0000000000..4d41fcdf60
--- /dev/null
+++ b/configs/s2ef/2M/mace/041801_specagn_withskip.yml
@@ -0,0 +1,87 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: SpeciesAgnosticResidualInteractionBlock
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/042401_srctgt_conv.yml b/configs/s2ef/2M/mace/042401_srctgt_conv.yml
new file mode 100644
index 0000000000..e9d57ea4e3
--- /dev/null
+++ b/configs/s2ef/2M/mace/042401_srctgt_conv.yml
@@ -0,0 +1,84 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: SpeciesAgnosticResidualInteractionBlock
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/042501_srctgt_conv_gradclip5.yml b/configs/s2ef/2M/mace/042501_srctgt_conv_gradclip5.yml
new file mode 100644
index 0000000000..68a71627af
--- /dev/null
+++ b/configs/s2ef/2M/mace/042501_srctgt_conv_gradclip5.yml
@@ -0,0 +1,84 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: SpeciesAgnosticResidualInteractionBlock
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 5
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/050101_srctgt_conv_h256.yml b/configs/s2ef/2M/mace/050101_srctgt_conv_h256.yml
new file mode 100644
index 0000000000..838494e362
--- /dev/null
+++ b/configs/s2ef/2M/mace/050101_srctgt_conv_h256.yml
@@ -0,0 +1,85 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: SpeciesAgnosticResidualInteractionBlock
+  contraction_type: v2
+  node_feats_down_irreps: 256x0e
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 5
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/050102_specdep_srctgt_conv.yml b/configs/s2ef/2M/mace/050102_specdep_srctgt_conv.yml
new file mode 100644
index 0000000000..e56a86dcdc
--- /dev/null
+++ b/configs/s2ef/2M/mace/050102_specdep_srctgt_conv.yml
@@ -0,0 +1,85 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: RealAgnosticResidualInteractionBlockV2
+  interaction_cls: RealAgnosticResidualInteractionBlockV2
+  contraction_type: v1
+  node_feats_down_irreps: 64x0e
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 8
+  eval_batch_size: 8
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/082201_edgegated.yml b/configs/s2ef/2M/mace/082201_edgegated.yml
new file mode 100644
index 0000000000..581c357c70
--- /dev/null
+++ b/configs/s2ef/2M/mace/082201_edgegated.yml
@@ -0,0 +1,85 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 256x0e + 256x1o + 256x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: EdgeGatedInteractionBlock
+  edge_gates_irreps: 128x0e
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 5
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/103101_edgegated_srctgt.yml b/configs/s2ef/2M/mace/103101_edgegated_srctgt.yml
new file mode 100644
index 0000000000..c4e9137873
--- /dev/null
+++ b/configs/s2ef/2M/mace/103101_edgegated_srctgt.yml
@@ -0,0 +1,87 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 256x0e + 256x1o + 256x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: EdgeGatedInteractionBlock
+  edge_gates_irreps: 128x0e
+  edge_gates_use_source_target_feats: True
+  node_feats_down_irreps: 64x0e
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 5
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/110501_edgegated_srctgt_h512.yml b/configs/s2ef/2M/mace/110501_edgegated_srctgt_h512.yml
new file mode 100644
index 0000000000..3fcfbcc601
--- /dev/null
+++ b/configs/s2ef/2M/mace/110501_edgegated_srctgt_h512.yml
@@ -0,0 +1,87 @@
+trainer: interaction_mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  train:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    normalize_labels: True
+    target_mean: -0.7554450631141663
+    target_std: 2.887317180633545
+    grad_target_mean: 0.0
+    grad_target_std: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: interaction_energy_mace
+  r_max: 12.0
+  max_neighbors: 30
+  avg_num_neighbors: 30
+  rbf: gaussian
+  num_rbf: 128
+  rbf_hidden_channels: 256
+  num_bessel: -1  # unused, but required.
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 56
+  hidden_irreps: 512x0e + 512x1o + 512x2e
+  MLP_irreps: 128x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20
+  interaction_cls_first: SpeciesAgnosticResidualInteractionBlock
+  interaction_cls: EdgeGatedInteractionBlock
+  edge_gates_irreps: 128x0e
+  edge_gates_use_source_target_feats: True
+  node_feats_down_irreps: 64x0e
+  contraction_type: v2
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+  direct_forces: True
+
+optim:
+  batch_size: 2
+  eval_batch_size: 2
+  load_balancing: atoms
+  eval_every: 2000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 5
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/2M/mace/base_0127.yml b/configs/s2ef/2M/mace/base_0127.yml
new file mode 100644
index 0000000000..8be2828f63
--- /dev/null
+++ b/configs/s2ef/2M/mace/base_0127.yml
@@ -0,0 +1,78 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  avg_num_neighbors: 35
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 3
+  num_elements: 83
+  hidden_irreps: 128x0e + 128x1o
+  MLP_irreps: 128x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  atomic_inter_shift: -0.75544
+  atomic_inter_scale: 2.88732
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 16
+  eval_batch_size: 16
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # Default in MACE (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  # was 1.e-2.
+  lr_initial: 5.e-3
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 40
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/020101.yml b/configs/s2ef/tiny/mace/020101.yml
new file mode 100644
index 0000000000..9708d938fe
--- /dev/null
+++ b/configs/s2ef/tiny/mace/020101.yml
@@ -0,0 +1,80 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/train/ads_energy
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/val/ads_energy
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  avg_num_neighbors: 35
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 83
+  hidden_irreps: 128x0e + 128x1o
+  MLP_irreps: 128x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.08374
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 64
+  eval_batch_size: 64
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/020102_exact_nbrs.yml b/configs/s2ef/tiny/mace/020102_exact_nbrs.yml
new file mode 100644
index 0000000000..953e08ae0a
--- /dev/null
+++ b/configs/s2ef/tiny/mace/020102_exact_nbrs.yml
@@ -0,0 +1,81 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/train/ads_energy
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/val/ads_energy
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: -1
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 83
+  hidden_irreps: 128x0e + 128x1o
+  MLP_irreps: 128x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.08374
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/020103_exact_nbrs_h256.yml b/configs/s2ef/tiny/mace/020103_exact_nbrs_h256.yml
new file mode 100644
index 0000000000..8eab45a03d
--- /dev/null
+++ b/configs/s2ef/tiny/mace/020103_exact_nbrs_h256.yml
@@ -0,0 +1,81 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/train/ads_energy/
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/val/ads_energy/
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: -1
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 83
+  hidden_irreps: 256x0e + 256x1o
+  MLP_irreps: 256x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.08374
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/020104_nbrs31_h256.yml b/configs/s2ef/tiny/mace/020104_nbrs31_h256.yml
new file mode 100644
index 0000000000..7e529fff5a
--- /dev/null
+++ b/configs/s2ef/tiny/mace/020104_nbrs31_h256.yml
@@ -0,0 +1,81 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/train/ads_energy/
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/val/ads_energy/
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: 31
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 83
+  hidden_irreps: 256x0e + 256x1o
+  MLP_irreps: 256x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.08374
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 100
+  energy_coefficient: 1
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  loss_energy: mae
+  loss_force: l2mae
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/020105_nbrs31_h256_f1000_e37.yml b/configs/s2ef/tiny/mace/020105_nbrs31_h256_f1000_e37.yml
new file mode 100644
index 0000000000..a3d3f3587f
--- /dev/null
+++ b/configs/s2ef/tiny/mace/020105_nbrs31_h256_f1000_e37.yml
@@ -0,0 +1,79 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/train/ads_energy/
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/val/ads_energy/
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: 31
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 83
+  hidden_irreps: 256x0e + 256x1o
+  MLP_irreps: 256x0e
+  # atomic_energies: ref energy per atomic number.
+  correlation: 3
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -0.08374
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/021801_total_energy.yml b/configs/s2ef/tiny/mace/021801_total_energy.yml
new file mode 100644
index 0000000000..a132c1b665
--- /dev/null
+++ b/configs/s2ef/tiny/mace/021801_total_energy.yml
@@ -0,0 +1,84 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/train/ads_energy/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/val/ads_energy/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: 31
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 83
+  hidden_irreps: 256x0e + 256x1o
+  MLP_irreps: 256x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20tiny
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -7.1012316731649495e-15
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462538421154022
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/022101_total_energy_zembedfix.yml b/configs/s2ef/tiny/mace/022101_total_energy_zembedfix.yml
new file mode 100644
index 0000000000..852b86aeb0
--- /dev/null
+++ b/configs/s2ef/tiny/mace/022101_total_energy_zembedfix.yml
@@ -0,0 +1,84 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/train/ads_energy/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/val/ads_energy/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: 31
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 5
+  hidden_irreps: 256x0e + 256x1o
+  MLP_irreps: 256x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20tiny
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -7.1012316731649495e-15
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462538421154022
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/022201_total_energy_h128_mlp16.yml b/configs/s2ef/tiny/mace/022201_total_energy_h128_mlp16.yml
new file mode 100644
index 0000000000..8d67fd0ff1
--- /dev/null
+++ b/configs/s2ef/tiny/mace/022201_total_energy_h128_mlp16.yml
@@ -0,0 +1,84 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: oc22_lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/train/ads_energy/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+  - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/v1_sid_str/val/ads_energy/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+  # - src: /checkpoint/bmwood/oc20_small/lmdbs/random1215212/test/ads_energy
+
+model:
+  name: scaleshift_mace
+  # Copied from https://github.com/ACEsuit/mace#training.
+  r_max: 6.0 # was originally 5.0
+  # Use `avg_num_neighbors` = 50 if r_max is 12.0, 35 if r_max is 6.0.
+  # This is because max_neighbors is set to 50. At r_max = 12, we get all nbrs.
+  # Passing -1 computes exact no. neighbors per system.
+  avg_num_neighbors: 31
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 5
+  hidden_irreps: 128x0e + 128x1o
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20tiny
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: -7.1012316731649495e-15
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.23462538421154022
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 64
+  eval_batch_size: 64
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/configs/s2ef/tiny/mace/022202_ilyes_weights.yml b/configs/s2ef/tiny/mace/022202_ilyes_weights.yml
new file mode 100644
index 0000000000..400f418830
--- /dev/null
+++ b/configs/s2ef/tiny/mace/022202_ilyes_weights.yml
@@ -0,0 +1,89 @@
+trainer: mace_trainer
+
+logger: wandb
+
+task:
+  dataset: lmdb  # for training on total energies.
+  primary_metric: forces_mae
+  train_on_free_atoms: False
+  eval_on_free_atoms: True
+  # OCP MACE models have "z2idx", which is missing in official MACE models.
+  # Official MACE models have "atomic_numbers", "r_max", "num_interactions",
+  # which are missing in OCP MACE models. Hence we set strict_load = False.
+  strict_load: False
+
+dataset:
+  - src: /checkpoint/bmwood/oc20_small/oc20_xxs_dataset/tot_e_lmdb/train/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # Switch off normalize_labels since scale_shift in the MACE model does the same thing.
+    normalize_labels: False
+    # target_mean: -0.7554450631141663
+    # target_std: 2.887317180633545
+    # grad_target_mean: 0.0
+    # grad_target_std: 2.887317180633545
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  - src: /checkpoint/bmwood/oc20_small/oc20_xs_dataset/tot_e_lmdb/val/
+    total_energy: True
+    oc20_ref: /checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl
+    # ase_nbrlist: True
+    # ase_cutoff: 8.0
+  # - src: /checkpoint/bmwood/oc20_small/oc20_xs_dataset/tot_e_lmdb/val/
+
+model:
+  name: scaleshift_mace
+  r_max: 8.0
+  max_neighbors: 500
+  avg_num_neighbors: 16.268049240112305
+  num_bessel: 8
+  num_polynomial_cutoff: 5
+  max_ell: 3
+  num_interactions: 2
+  num_elements: 5
+  hidden_irreps: 256x0e + 256x1o + 256x2e
+  MLP_irreps: 16x0e
+  correlation: 3
+  # atomic_energies: ref energy per atomic number.
+  atomic_energies: oc20xxs
+  # torch.div(energies, natoms).mean()
+  atomic_inter_shift: 0.0046684232644386295
+  # torch.sqrt(torch.mean(torch.square(forces)))
+  atomic_inter_scale: 0.37896862626075745
+  # OCP flags.
+  regress_forces: True
+  otf_graph: True
+  use_pbc: True
+
+optim:
+  batch_size: 32
+  eval_batch_size: 32
+  load_balancing: atoms
+  eval_every: 1000
+  num_workers: 2
+  # default from (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L234)
+  lr_initial: 1.e-2
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  # In the MACE code, the RLROP scheduler is run every epoch even if validation loss is stale.
+  # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/mace/tools/train.py#L149
+  # In the OCP code, we step the scheduler only when validation loss is recomputed.
+  # https://github.com/ACEsuit/mace-ocp/blob/main/ocpmodels/trainers/forces_trainer.py#L412
+  # And so we should probably keep a lower patience. Using default from OCP for now.
+  patience: 3
+  max_epochs: 1000
+  # TODO(@abhshkdz): implement the swa scheduler, i.e. switch to high energy coefficient and low force coefficient after loss plateaus with initial coefficients.
+  force_coefficient: 1000
+  energy_coefficient: 37
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L285
+  ema_decay: 0.99
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L321
+  clip_grad_norm: 10
+  # default from https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L240.
+  weight_decay: 5.e-7
+
+slurm:
+  constraint: "volta32gb"
diff --git a/ocpmodels/datasets/oc22_lmdb_dataset.py b/ocpmodels/datasets/oc22_lmdb_dataset.py
index 575eedc86b..57b8e020bb 100644
--- a/ocpmodels/datasets/oc22_lmdb_dataset.py
+++ b/ocpmodels/datasets/oc22_lmdb_dataset.py
@@ -13,6 +13,8 @@
 import warnings
 from pathlib import Path
 
+import ase.neighborlist
+
 import lmdb
 import numpy as np
 import torch
@@ -184,6 +186,34 @@ def __getitem__(self, idx):
         if "distances" in data_object:
             del data_object.distances
 
+        if self.config.get("ase_nbrlist", False):
+            sender, receiver, unit_shifts = ase.neighborlist.primitive_neighbor_list(
+                quantities="ijS",
+                pbc=(True, True, True),
+                cell=data_object.cell[0].numpy(),
+                positions=data_object.pos.numpy(),
+                cutoff=self.config.get("ase_cutoff", 6.0),
+                self_interaction=True,  # we want edges from atom to itself in different periodic images
+                use_scaled_positions=False,  # positions are not scaled positions
+            )
+
+            # Eliminate self-edges that don't cross periodic boundaries
+            true_self_edge = sender == receiver
+            true_self_edge &= np.all(unit_shifts == 0, axis=1)
+            keep_edge = ~true_self_edge
+
+            # Note: after eliminating self-edges, it can be that no edges remain in this system
+            sender = sender[keep_edge]
+            receiver = receiver[keep_edge]
+            unit_shifts = unit_shifts[keep_edge]
+
+            data_object.edge_index = torch.from_numpy(np.stack((sender, receiver)))  # [2, n_edges]
+            data_object.unit_shifts = torch.from_numpy(unit_shifts)
+
+            # From the docs: With the shift vector S, the distances D between atoms can be computed from
+            # D = positions[j]-positions[i]+S.dot(cell)
+            data_object.shifts = torch.from_numpy(np.dot(unit_shifts, data_object.cell[0])).to(torch.float32)  # [n_edges, 3]
+
         return data_object
 
     def connect_db(self, lmdb_path=None):
diff --git a/ocpmodels/models/mace/blocks.py b/ocpmodels/models/mace/blocks.py
new file mode 100644
index 0000000000..3411906ea4
--- /dev/null
+++ b/ocpmodels/models/mace/blocks.py
@@ -0,0 +1,899 @@
+from abc import ABC, abstractmethod
+from typing import Callable, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional
+from e3nn import nn, o3
+from e3nn.util.jit import compile_mode
+
+from .mace_core.blocks import InteractionBlock
+from .mace_core.irreps_tools import (
+    lifted_skip,
+    linear_out_irreps,
+    reshape_irreps,
+    tp_out_irreps_with_instructions,
+)
+from .mace_core.scatter import scatter_sum
+from .mace_core.symmetric_contraction import SymmetricContraction
+from .radial import BesselBasis, GaussianBasis, PolynomialCutoff
+
+
+class LinearNodeEmbeddingBlock(torch.nn.Module):
+    def __init__(self, irreps_in: o3.Irreps, irreps_out: o3.Irreps):
+        super().__init__()
+        self.linear = o3.Linear(irreps_in=irreps_in, irreps_out=irreps_out)
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,  # [n_nodes, irreps]
+    ):
+        return self.linear(node_attrs)
+
+
+class LinearReadoutBlock(torch.nn.Module):
+    def __init__(self, irreps_in: o3.Irreps):
+        super().__init__()
+        self.linear = o3.Linear(
+            irreps_in=irreps_in, irreps_out=o3.Irreps("0e")
+        )
+
+    def forward(
+        self, x: torch.Tensor
+    ) -> torch.Tensor:  # [n_nodes, irreps]  # [..., ]
+        return self.linear(x)  # [n_nodes, 1]
+
+
+class NonLinearReadoutBlock(torch.nn.Module):
+    def __init__(
+        self,
+        irreps_in: o3.Irreps,
+        MLP_irreps: o3.Irreps,
+        gate: Optional[Callable],
+    ):
+        super().__init__()
+        self.hidden_irreps = MLP_irreps
+        self.linear_1 = o3.Linear(
+            irreps_in=irreps_in, irreps_out=self.hidden_irreps
+        )
+        self.non_linearity = nn.Activation(
+            irreps_in=self.hidden_irreps, acts=[gate]
+        )
+        self.linear_2 = o3.Linear(
+            irreps_in=self.hidden_irreps, irreps_out=o3.Irreps("0e")
+        )
+
+    def forward(
+        self, x: torch.Tensor
+    ) -> torch.Tensor:  # [n_nodes, irreps]  # [..., ]
+        x = self.non_linearity(self.linear_1(x))
+        return self.linear_2(x)  # [n_nodes, 1]
+
+
+class AtomicEnergiesBlock(torch.nn.Module):
+    atomic_energies: torch.Tensor
+
+    def __init__(self, atomic_energies: Union[np.ndarray, torch.Tensor]):
+        super().__init__()
+        assert len(atomic_energies.shape) == 1
+
+        self.register_buffer(
+            "atomic_energies",
+            torch.tensor(atomic_energies, dtype=torch.get_default_dtype()),
+        )  # [n_elements, ]
+
+    def forward(
+        self, x: torch.Tensor  # one-hot of elements [..., n_elements]
+    ) -> torch.Tensor:  # [..., ]
+        return torch.matmul(x, self.atomic_energies)
+
+    def __repr__(self):
+        formatted_energies = ", ".join(
+            [f"{x:.4f}" for x in self.atomic_energies]
+        )
+        return f"{self.__class__.__name__}(energies=[{formatted_energies}])"
+
+
+class RadialEmbeddingBlock(torch.nn.Module):
+    def __init__(
+        self,
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        rbf: str = "bessel",
+    ):
+        super().__init__()
+        self.rbf = rbf
+        if rbf == "bessel":
+            self.bessel_fn = BesselBasis(r_max=r_max, num_basis=num_bessel)
+        elif rbf == "gaussian":
+            self.rbf_fn = GaussianBasis(r_max=r_max, num_gaussians=num_bessel)
+        self.cutoff_fn = PolynomialCutoff(r_max=r_max, p=num_polynomial_cutoff)
+        self.out_dim = num_bessel
+
+    def forward(
+        self,
+        edge_lengths: torch.Tensor,  # [n_edges, 1]
+    ):
+        if self.rbf == "bessel":
+            bessel = self.bessel_fn(edge_lengths)  # [n_edges, n_basis]
+            cutoff = self.cutoff_fn(edge_lengths)  # [n_edges, 1]
+            return bessel * cutoff  # [n_edges, n_basis]
+        elif self.rbf == "gaussian":
+            gaussian = self.rbf_fn(edge_lengths)  # [n_edges, n_basis]
+            cutoff = self.cutoff_fn(edge_lengths)  # [n_edges, 1]
+            return gaussian * cutoff  # [n_edges, n_basis]
+
+
+nonlinearities = {1: torch.nn.SiLU(), -1: torch.nn.Tanh()}
+
+
+class TensorProductWeightsBlock(torch.nn.Module):
+    def __init__(
+        self, num_elements: int, num_edge_feats: int, num_feats_out: int
+    ):
+        super().__init__()
+
+        weights = torch.empty(
+            (num_elements, num_edge_feats, num_feats_out),
+            dtype=torch.get_default_dtype(),
+        )
+        torch.nn.init.xavier_uniform_(weights)
+        self.weights = torch.nn.Parameter(weights)
+
+    def forward(
+        self,
+        sender_or_receiver_node_attrs: torch.Tensor,  # assumes that the node attributes are one-hot encoded
+        edge_feats: torch.Tensor,
+    ):
+        return torch.einsum(
+            "be, ba, aek -> bk",
+            edge_feats,
+            sender_or_receiver_node_attrs,
+            self.weights,
+        )
+
+    def __repr__(self):
+        return (
+            f'{self.__class__.__name__}(shape=({", ".join(str(s) for s in self.weights.shape)}), '
+            f"weights={np.prod(self.weights.shape)})"
+        )
+
+
+class ResidualElementDependentInteractionBlock(InteractionBlock):
+    def _setup(self) -> None:
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps, self.edge_attrs_irreps, self.target_irreps
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+        self.conv_tp_weights = TensorProductWeightsBlock(
+            num_elements=self.node_attrs_irreps.num_irreps,
+            num_edge_feats=self.edge_feats_irreps.num_irreps,
+            num_feats_out=self.conv_tp.weight_numel,
+        )
+
+        # Linear
+        irreps_mid = irreps_mid.simplify()
+        self.irreps_out = linear_out_irreps(irreps_mid, self.target_irreps)
+        self.irreps_out = self.irreps_out.simplify()
+        self.linear = o3.Linear(
+            irreps_mid,
+            self.irreps_out,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        # Selector TensorProduct
+        self.skip_tp = o3.FullyConnectedTensorProduct(
+            self.node_feats_irreps, self.node_attrs_irreps, self.irreps_out
+        )
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> torch.Tensor:
+        sender, receiver = edge_index
+        num_nodes = node_feats.shape[0]
+        sc = self.skip_tp(node_feats, node_attrs)
+        node_feats = self.linear_up(node_feats)
+        tp_weights = self.conv_tp_weights(node_attrs[sender], edge_feats)
+        mji = self.conv_tp(
+            node_feats[sender], edge_attrs, tp_weights
+        )  # [n_edges, irreps]
+        message = scatter_sum(
+            src=mji, index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, irreps]
+        message = self.linear(message) / self.avg_num_neighbors
+        return message + sc  # [n_nodes, irreps]
+
+
+class AgnosticNonlinearInteractionBlock(InteractionBlock):
+    def _setup(self) -> None:
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps, self.edge_attrs_irreps, self.target_irreps
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+
+        # Convolution weights
+        input_dim = self.edge_feats_irreps.num_irreps
+        self.conv_tp_weights = nn.FullyConnectedNet(
+            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel],
+            torch.nn.SiLU(),
+        )
+
+        # Linear
+        irreps_mid = irreps_mid.simplify()
+        self.irreps_out = linear_out_irreps(irreps_mid, self.target_irreps)
+        self.irreps_out = self.irreps_out.simplify()
+        self.linear = o3.Linear(
+            irreps_mid,
+            self.irreps_out,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        # Selector TensorProduct
+        self.skip_tp = o3.FullyConnectedTensorProduct(
+            self.irreps_out, self.node_attrs_irreps, self.irreps_out
+        )
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> torch.Tensor:
+        sender, receiver = edge_index
+        num_nodes = node_feats.shape[0]
+        tp_weights = self.conv_tp_weights(edge_feats)
+        node_feats = self.linear_up(node_feats)
+        mji = self.conv_tp(
+            node_feats[sender], edge_attrs, tp_weights
+        )  # [n_edges, irreps]
+        message = scatter_sum(
+            src=mji, index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, irreps]
+        message = self.linear(message) / self.avg_num_neighbors
+        message = self.skip_tp(message, node_attrs)
+        return message  # [n_nodes, irreps]
+
+
+class AgnosticResidualNonlinearInteractionBlock(InteractionBlock):
+    def _setup(self) -> None:
+        # First linear
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps, self.edge_attrs_irreps, self.target_irreps
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+
+        # Convolution weights
+        input_dim = self.edge_feats_irreps.num_irreps
+        self.conv_tp_weights = nn.FullyConnectedNet(
+            [input_dim] + 3 * [64] + [self.conv_tp.weight_numel],
+            torch.nn.SiLU(),
+        )
+
+        # Linear
+        irreps_mid = irreps_mid.simplify()
+        self.irreps_out = linear_out_irreps(irreps_mid, self.target_irreps)
+        self.irreps_out = self.irreps_out.simplify()
+        self.linear = o3.Linear(
+            irreps_mid,
+            self.irreps_out,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        # Selector TensorProduct
+        self.skip_tp = o3.FullyConnectedTensorProduct(
+            self.node_feats_irreps, self.node_attrs_irreps, self.irreps_out
+        )
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> torch.Tensor:
+        sender, receiver = edge_index
+        num_nodes = node_feats.shape[0]
+        sc = self.skip_tp(node_feats, node_attrs)
+        node_feats = self.linear_up(node_feats)
+        tp_weights = self.conv_tp_weights(edge_feats)
+        mji = self.conv_tp(
+            node_feats[sender], edge_attrs, tp_weights
+        )  # [n_edges, irreps]
+        message = scatter_sum(
+            src=mji, index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, irreps]
+        message = self.linear(message) / self.avg_num_neighbors
+        message = message + sc
+        return message  # [n_nodes, irreps]
+
+
+class ScaleShiftBlock(torch.nn.Module):
+    def __init__(self, scale: float, shift: float):
+        super().__init__()
+        self.register_buffer(
+            "scale", torch.tensor(scale, dtype=torch.get_default_dtype())
+        )
+        self.register_buffer(
+            "shift", torch.tensor(shift, dtype=torch.get_default_dtype())
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.scale * x + self.shift
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(scale={self.scale:.6f}, shift={self.shift:.6f})"
+
+
+class ForceBlock(torch.nn.Module):
+    def __init__(self, hidden_irreps):
+        super().__init__()
+
+        # TODO(@abhshkdz): is this assertion needed?
+        self.hidden_irreps = hidden_irreps
+        assert hidden_irreps[0].dim * 3 == hidden_irreps[1].dim
+
+        l0_h = hidden_irreps[0].mul
+        l1_h = hidden_irreps[1].mul
+
+        self.output_network = torch.nn.ModuleList(
+            [
+                GatedEquivariantBlock(l0_h, l1_h, l1_h // 2),
+                GatedEquivariantBlock(l1_h // 2, l1_h // 2, 1),
+            ]
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for layer in self.output_network:
+            layer.reset_parameters()
+
+    def forward(self, node_feats):
+        # split node_feats into scalar and vector components.
+        x = node_feats[:, : self.hidden_irreps[0].dim]
+        vec = node_feats[
+            :,
+            self.hidden_irreps[0].dim : self.hidden_irreps[0].dim
+            + self.hidden_irreps[1].dim,
+        ]
+
+        vec = vec.reshape(vec.shape[0], self.hidden_irreps[1].mul, 3)
+        vec = vec.transpose(1, 2)
+
+        # pass it through the gated equivariant blocks.
+        for layer in self.output_network:
+            x, vec = layer(x, vec)
+
+        # return the vector components.
+        return vec.squeeze()
+
+
+# Implementation based on TorchMD-Net
+# https://github.com/Open-Catalyst-Project/ocp/blob/main/ocpmodels/models/painn/painn.py#L607
+class GatedEquivariantBlock(torch.nn.Module):
+    """Gated Equivariant Block as defined in Schütt et al. (2021):
+    Equivariant message passing for the prediction of tensorial properties and molecular spectra
+    """
+
+    def __init__(
+        self,
+        l0_channels,
+        l1_channels,
+        out_channels,
+    ):
+        super(GatedEquivariantBlock, self).__init__()
+        self.l0_channels = l0_channels
+        self.l1_channels = l1_channels
+        self.out_channels = out_channels
+
+        self.vec1_proj = torch.nn.Linear(l1_channels, l1_channels, bias=False)
+        self.vec2_proj = torch.nn.Linear(l1_channels, out_channels, bias=False)
+
+        self.update_net = torch.nn.Sequential(
+            torch.nn.Linear(
+                l0_channels + l1_channels,
+                l1_channels,
+            ),
+            torch.nn.SiLU(),
+            torch.nn.Linear(l1_channels, out_channels * 2),
+        )
+
+        self.act = torch.nn.SiLU()
+
+    def reset_parameters(self):
+        torch.nn.init.xavier_uniform_(self.vec1_proj.weight)
+        torch.nn.init.xavier_uniform_(self.vec2_proj.weight)
+        torch.nn.init.xavier_uniform_(self.update_net[0].weight)
+        self.update_net[0].bias.data.fill_(0)
+        torch.nn.init.xavier_uniform_(self.update_net[2].weight)
+        self.update_net[2].bias.data.fill_(0)
+
+    def forward(self, x, v):
+        # x is [num_nodes x l0_hidden_channels]
+        # v is [num_nodes x 3 x l1_hidden_channels]
+        vec1 = torch.norm(self.vec1_proj(v), dim=-2)
+        vec2 = self.vec2_proj(v)
+
+        x = torch.cat([x, vec1], dim=-1)
+        x, v = torch.split(self.update_net(x), self.out_channels, dim=-1)
+        v = v.unsqueeze(1) * vec2
+
+        x = self.act(x)
+        return x, v
+
+
+@compile_mode("script")
+class SpeciesAgnosticResidualInteractionBlock(InteractionBlock):
+    def _setup(self) -> None:
+        # First linear
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            self.target_irreps,
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+
+        # Convolution weights
+        self.linear_down = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_down_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        input_dim = (
+            self.edge_feats_irreps.num_irreps
+            + 2 * self.node_feats_down_irreps.num_irreps
+        )
+        self.conv_tp_weights = nn.FullyConnectedNet(
+            [input_dim]
+            + 3 * [self.rbf_hidden_channels]
+            + [self.conv_tp.weight_numel],
+            torch.nn.functional.silu,
+        )
+
+        # Linear
+        irreps_mid = irreps_mid.simplify()
+        self.irreps_out = self.target_irreps
+        self.linear = o3.Linear(
+            irreps_mid,
+            self.irreps_out,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        self.reshape = reshape_irreps(self.irreps_out)
+
+        # Skip connection.
+        self.skip_linear = o3.Linear(
+            self.node_feats_irreps, self.hidden_irreps
+        )
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> Tuple[torch.Tensor, None]:
+        sender = edge_index[0]
+        receiver = edge_index[1]
+        num_nodes = node_feats.shape[0]
+        sc = self.skip_linear(node_feats)
+        node_feats_up = self.linear_up(node_feats)
+        node_feats_down = self.linear_down(node_feats)
+        augmented_edge_feats = torch.cat(
+            [
+                edge_feats,
+                node_feats_down[sender],
+                node_feats_down[receiver],
+            ],
+            dim=-1,
+        )
+        tp_weights = self.conv_tp_weights(augmented_edge_feats)
+        mji = self.conv_tp(
+            node_feats_up[sender], edge_attrs, tp_weights
+        )  # [n_edges, irreps]
+        message = scatter_sum(
+            src=mji, index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, irreps]
+        message = self.linear(message) / self.avg_num_neighbors
+        return (
+            self.reshape(message),
+            sc,
+        )  # [n_nodes, channels, (lmax + 1)**2]
+
+
+@compile_mode("script")
+class RealAgnosticResidualInteractionBlockV2(InteractionBlock):
+    def _setup(self) -> None:
+
+        # First linear
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            self.target_irreps,
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+
+        # Convolution weights
+        self.linear_down = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_down_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        input_dim = (
+            self.edge_feats_irreps.num_irreps
+            + 2 * self.node_feats_down_irreps.num_irreps
+        )
+        self.conv_tp_weights = nn.FullyConnectedNet(
+            [input_dim]
+            + 3 * [self.rbf_hidden_channels]
+            + [self.conv_tp.weight_numel],
+            torch.nn.functional.silu,
+        )
+
+        # Linear
+        irreps_mid = irreps_mid.simplify()
+        self.irreps_out = self.target_irreps
+        self.linear = o3.Linear(
+            irreps_mid,
+            self.irreps_out,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        # Selector TensorProduct
+        self.skip_tp = o3.FullyConnectedTensorProduct(
+            self.node_feats_irreps, self.node_attrs_irreps, self.hidden_irreps
+        )
+        self.reshape = reshape_irreps(self.irreps_out)
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        sender = edge_index[0]
+        receiver = edge_index[1]
+        num_nodes = node_feats.shape[0]
+        sc = self.skip_tp(node_feats, node_attrs)
+        node_feats_up = self.linear_up(node_feats)
+        node_feats_down = self.linear_down(node_feats)
+        augmented_edge_feats = torch.cat(
+            [
+                edge_feats,
+                node_feats_down[sender],
+                node_feats_down[receiver],
+            ],
+            dim=-1,
+        )
+        tp_weights = self.conv_tp_weights(augmented_edge_feats)
+        mji = self.conv_tp(
+            node_feats_up[sender], edge_attrs, tp_weights
+        )  # [n_edges, irreps]
+        message = scatter_sum(
+            src=mji, index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, irreps]
+        message = self.linear(message) / self.avg_num_neighbors
+        return (
+            self.reshape(message),
+            sc,
+        )  # [n_nodes, channels, (lmax + 1)**2]
+
+
+@compile_mode("script")
+class IdentityResidualInteractionBlock(InteractionBlock):
+    def _setup(self) -> None:
+        # First linear
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            self.target_irreps,
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+
+        # Convolution weights
+        input_dim = self.edge_feats_irreps.num_irreps
+        self.conv_tp_weights = nn.FullyConnectedNet(
+            [input_dim]
+            + 3 * [self.rbf_hidden_channels]
+            + [self.conv_tp.weight_numel],
+            torch.nn.functional.silu,
+        )
+
+        # Linear
+        irreps_mid = irreps_mid.simplify()
+        self.irreps_out = self.target_irreps
+        self.linear = o3.Linear(
+            irreps_mid,
+            self.irreps_out,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        self.skip = lifted_skip(self.node_feats_irreps, self.hidden_irreps)
+        self.reshape = reshape_irreps(self.irreps_out)
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> Tuple[torch.Tensor, None]:
+        sender = edge_index[0]
+        receiver = edge_index[1]
+        num_nodes = node_feats.shape[0]
+        sc = self.skip(node_feats)
+        node_feats = self.linear_up(node_feats)
+        tp_weights = self.conv_tp_weights(edge_feats)
+        mji = self.conv_tp(
+            node_feats[sender], edge_attrs, tp_weights
+        )  # [n_edges, irreps]
+        message = scatter_sum(
+            src=mji, index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, irreps]
+        message = self.linear(message) / self.avg_num_neighbors
+        return (
+            self.reshape(message),
+            sc,
+        )  # [n_nodes, channels, (lmax + 1)**2]
+
+
+@compile_mode("script")
+class EdgeGatedInteractionBlock(InteractionBlock):
+    def _setup(self) -> None:
+
+        # First linear
+        self.linear_up = o3.Linear(
+            self.node_feats_irreps,
+            self.node_feats_irreps,
+            internal_weights=True,
+            shared_weights=True,
+        )
+
+        # TensorProduct
+        irreps_mid, instructions = tp_out_irreps_with_instructions(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            self.target_irreps,
+        )
+        self.conv_tp = o3.TensorProduct(
+            self.node_feats_irreps,
+            self.edge_attrs_irreps,
+            irreps_mid,
+            instructions=instructions,
+            shared_weights=False,
+            internal_weights=False,
+        )
+        irreps_mid = irreps_mid.simplify()
+        self.message_dim = irreps_mid.dim
+        self.irreps_out = self.target_irreps
+
+        # Convolution weights
+        self.num_convs = self.num_gates if self.multi_conv else 1
+
+        if self.use_source_target_edge_feats:
+            self.linear_down = o3.Linear(
+                self.node_feats_irreps,
+                self.node_feats_down_irreps,
+                internal_weights=True,
+                shared_weights=True,
+            )
+            input_dim = (
+                self.edge_feats_irreps.num_irreps
+                + 2 * self.node_feats_down_irreps.num_irreps
+            )
+            self.conv_tp_weights = nn.FullyConnectedNet(
+                [input_dim]
+                + 3 * [self.rbf_hidden_channels]
+                + [self.conv_tp.weight_numel * self.num_convs],
+                torch.nn.functional.silu,
+            )
+        else:
+            input_dim = self.edge_feats_irreps.num_irreps
+            self.conv_tp_weights = nn.FullyConnectedNet(
+                [input_dim]
+                + 3 * [self.rbf_hidden_channels]
+                + [self.conv_tp.weight_numel * self.num_convs],
+                torch.nn.functional.silu,
+            )
+
+        # Edge gating operations
+        if self.exponential:
+            # self.scale = nn.FullyConnectedNet(
+            #     [self.edge_gates_irreps.dim] + 3 * [self.rbf_hidden_channels] + [1],
+            #     torch.nn.functional.silu,
+            # )
+            self.scale = None
+        self.alpha = o3.Linear(self.node_feats_irreps, self.edge_gates_irreps)
+        self.gamma = o3.Linear(
+            irreps_mid * self.num_convs, self.edge_gates_irreps
+        )
+        self.mix = o3.FullyConnectedTensorProduct(
+            self.edge_gates_irreps,
+            self.edge_gates_irreps,
+            f"{self.num_gates}x0e",
+        )
+        self.project = o3.Linear(irreps_mid * self.num_gates, self.irreps_out)
+
+        self.skip = lifted_skip(self.node_feats_irreps, self.hidden_irreps)
+        self.reshape = reshape_irreps(self.irreps_out)
+
+    def forward(
+        self,
+        node_attrs: torch.Tensor,
+        node_feats: torch.Tensor,
+        edge_attrs: torch.Tensor,
+        edge_feats: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        sender = edge_index[0]
+        receiver = edge_index[1]
+        num_nodes = node_feats.shape[0]
+        num_edges = edge_feats.shape[0]
+        sc = self.skip(node_feats)
+        node_feats_up = self.linear_up(node_feats)
+        if self.use_source_target_edge_feats:
+            node_feats_down = self.linear_down(node_feats)
+            augmented_edge_feats = torch.cat(
+                [
+                    edge_feats,
+                    node_feats_down[sender],
+                    node_feats_down[receiver],
+                ],
+                dim=-1,
+            )
+            tp_weights = self.conv_tp_weights(augmented_edge_feats).reshape(
+                num_edges, self.num_convs, self.conv_tp.weight_numel
+            )
+        else:
+            tp_weights = self.conv_tp_weights(edge_feats).reshape(
+                num_edges, self.num_convs, self.conv_tp.weight_numel
+            )
+
+        mji = self.conv_tp(
+            node_feats_up[sender][:, None, :],
+            edge_attrs[:, None, :],
+            tp_weights,
+        )  # [n_edges, n_convs, irreps]
+
+        # Gate edges
+        alphas, gammas = (
+            self.alpha(node_feats_up[receiver]),
+            self.gamma(
+                mji.reshape(num_edges, self.message_dim * self.num_convs)
+            ),
+        )
+        gate = self.mix(alphas, gammas)  # [n_edges, n_gates]
+        if self.exponential:
+            # local_scale = torch.nn.functional.softplus(
+            #     self.scale(alphas)
+            # )
+            # gate = (gate / local_scale).exp()
+            #
+            # Comment(@abhshkdz): The scatter_sum errors out because all_gate
+            # has infs. The temperature scaling function is likely causing this.
+            # Let's train with a fixed temperature for now.
+            gate = gate.exp()
+
+        all_gate = scatter_sum(
+            src=gate, index=receiver, dim=0, dim_size=num_nodes
+        )
+        gate = torch.nan_to_num(gate / all_gate[receiver]).unsqueeze(
+            -1
+        )  # [n_edges, n_gates, 1]
+
+        message = scatter_sum(
+            src=(gate * mji), index=receiver, dim=0, dim_size=num_nodes
+        )  # [n_nodes, n_gates, irreps]
+        message = message.reshape(
+            num_nodes, self.message_dim * self.num_gates
+        )  # [n_nodes, irreps * n_gates]
+        message = self.project(message)  # [n_nodes, irreps_out]
+        return (
+            self.reshape(message),
+            sc,
+        )
diff --git a/ocpmodels/models/mace/models.py b/ocpmodels/models/mace/models.py
new file mode 100644
index 0000000000..e63b6eb0bd
--- /dev/null
+++ b/ocpmodels/models/mace/models.py
@@ -0,0 +1,668 @@
+from typing import Any, Callable, Dict, List, Optional, Type
+
+import numpy as np
+import torch
+from e3nn import o3
+
+from ocpmodels.common.registry import registry
+from ocpmodels.common.utils import conditional_grad
+from ocpmodels.models.base import BaseModel
+
+from .blocks import (
+    AtomicEnergiesBlock,
+    EdgeGatedInteractionBlock,
+    ForceBlock,
+    IdentityResidualInteractionBlock,
+    LinearNodeEmbeddingBlock,
+    LinearReadoutBlock,
+    NonLinearReadoutBlock,
+    RadialEmbeddingBlock,
+    RealAgnosticResidualInteractionBlockV2,
+    ScaleShiftBlock,
+    SpeciesAgnosticResidualInteractionBlock,
+)
+from .mace_core.blocks import (
+    AgnosticInteractionBlock,
+    EquivariantProductBasisBlock,
+    RealAgnosticResidualInteractionBlock,
+)
+from .mace_core.scatter import scatter_sum
+from .utils import compute_forces
+
+
+@registry.register_model("mace")
+class MACE(BaseModel):
+    def __init__(
+        self,
+        # Unused legacy OCP params.
+        num_atoms: Optional[int],
+        bond_feat_dim: int,
+        num_targets: int,
+        #
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        max_ell: int,
+        num_interactions: int,
+        num_elements: int,
+        hidden_irreps: str,
+        MLP_irreps: str,
+        avg_num_neighbors: float,
+        correlation: int,
+        # Defaults from OCP / https://github.com/ACEsuit/mace/blob/main/scripts/run_train.py
+        gate=torch.nn.functional.silu,
+        # per-element energy references currently initialized to 0s.
+        atomic_energies=str,
+        interaction_cls=RealAgnosticResidualInteractionBlock,
+        interaction_cls_first=RealAgnosticResidualInteractionBlock,
+        max_neighbors: int = 500,
+        otf_graph: bool = True,
+        use_pbc: bool = True,
+        regress_forces: bool = True,
+        # support for gaussian basis.
+        rbf: str = "bessel",
+        num_rbf: int = 8,
+        rbf_hidden_channels: int = 64,
+        direct_forces: bool = False,
+        # support for species-agonstic contraction.
+        contraction_type: str = "v1",
+        # source and target feature size when concatenating for edge conv.
+        node_feats_down_irreps: str = "64x0e",
+        # parameters for edge gating.
+        edge_gates_irreps: str = "0e",
+        num_gates: int = 4,
+        multi_conv: bool = False,
+        exponential: bool = True,
+        edge_gates_use_source_target_feats: bool = False,
+    ):
+        super().__init__()
+        self.cutoff = self.r_max = r_max
+        self.avg_num_neighbors = avg_num_neighbors
+        self.max_neighbors = max_neighbors
+        self.otf_graph = otf_graph
+        self.use_pbc = use_pbc
+        self.regress_forces = regress_forces
+        self.direct_forces = direct_forces
+
+        # YAML loads them as strings, initialize them as o3.Irreps.
+        hidden_irreps = o3.Irreps(hidden_irreps)
+        MLP_irreps = o3.Irreps(MLP_irreps)
+
+        # Embedding
+        node_attr_irreps = o3.Irreps([(num_elements, (0, 1))])
+        node_feats_irreps = o3.Irreps(
+            [(hidden_irreps.count(o3.Irrep(0, 1)), (0, 1))]
+        )
+        self.node_embedding = LinearNodeEmbeddingBlock(
+            irreps_in=node_attr_irreps, irreps_out=node_feats_irreps
+        )
+        if rbf == "bessel":
+            self.radial_embedding = RadialEmbeddingBlock(
+                r_max=r_max,
+                num_bessel=num_bessel,
+                num_polynomial_cutoff=num_polynomial_cutoff,
+                rbf="bessel",
+            )
+        elif rbf == "gaussian":
+            self.radial_embedding = RadialEmbeddingBlock(
+                r_max=r_max,
+                num_bessel=num_rbf,
+                num_polynomial_cutoff=num_polynomial_cutoff,
+                rbf="gaussian",
+            )
+        edge_feats_irreps = o3.Irreps(f"{self.radial_embedding.out_dim}x0e")
+
+        sh_irreps = o3.Irreps.spherical_harmonics(max_ell)
+        num_features = hidden_irreps.count(o3.Irrep(0, 1))
+        interaction_irreps = (sh_irreps * num_features).sort()[0].simplify()
+        self.spherical_harmonics = o3.SphericalHarmonics(
+            sh_irreps, normalize=True, normalization="component"
+        )
+
+        edge_gates_irreps = o3.Irreps(edge_gates_irreps)
+
+        # Node / e0 block
+        # fmt: off
+        if atomic_energies == "oc20tiny":
+            atomic_energies = np.array([-0.5789446234902406, 0.0, 0.0, 0.0, 0.0, -0.5789446234902277, 0.0, -0.28947231174511384, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -4.631556987921935, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -4.631556987921935, 0.0, 0.0, 0.0])
+        elif atomic_energies == "oc20xxs":
+            atomic_energies = np.array([-3.8315212119598576, 0.0, 0.0, 0.0, 0.0, -7.902387218067884, -7.21000990963105, -6.879033386904713, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6946341972298161, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
+        elif atomic_energies == "oc20":
+            atomic_energies = np.array([-3.39973399802097, 0.0, 0.0, 0.0, -6.081884367921049, -8.293556029600811, -8.41745373621903, -4.92136266857127, 0.0, 0.0, -1.5412059197319175, 0.0, -3.6178325363353707, -5.3092430681268095, -5.291202563999228, -4.612401891658523, -3.081651431825209, 0.0, -1.5317668260051613, -2.4421854893794492, -6.728220402149601, -7.753628720142588, -8.380304712792668, -8.438050121320343, -7.841092458071865, -7.184260107389338, -6.23486296929759, -4.888130178461133, -2.962084764498329, -0.6657573646576105, -2.573145196892305, -4.067466705751315, -4.493942029750251, -3.811815951016779, 0.0, 0.0, -1.3488687481117176, -2.347882569119374, -7.026159902128673, -8.536582042496924, -9.587336073043936, -9.83366289354133, -9.268529057192797, -8.380115555600332, -6.794135382710533, -4.787344383569342, -1.8919937285040638, -0.20987260293014098, -2.0646188038840636, -3.3847372997774205, -3.648226709090113, -3.0333477411720318, 0.0, 0.0, -1.4964734007514156, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -9.921176307016488, -11.142905408621559, -11.666772123514626, -11.047390248000076, -10.13349975986731, -8.286331766183878, -5.775058392981709, -2.6918083582766843, 0.28829374995119617, -1.808141562387085, -3.0229565921247694, -3.329866054812059])
+        else:
+            atomic_energies = np.array([0. for i in range(83)])
+        # fmt: on
+
+        self.nz_z = np.nonzero(atomic_energies)[0]
+        assert len(self.nz_z) == num_elements
+
+        self.register_buffer(
+            "z2idx",
+            torch.zeros(120, dtype=torch.int64).fill_(-1),
+        )
+        for i, z in enumerate(self.nz_z + 1):
+            self.z2idx[z] = i
+
+        self.atomic_energies_fn = AtomicEnergiesBlock(
+            atomic_energies[self.nz_z]
+        )
+
+        # Interactions and readout
+        inter = interaction_cls_first(
+            node_attrs_irreps=node_attr_irreps,
+            node_feats_irreps=node_feats_irreps,
+            edge_attrs_irreps=sh_irreps,
+            edge_feats_irreps=edge_feats_irreps,
+            target_irreps=interaction_irreps,
+            hidden_irreps=hidden_irreps,
+            # MACE considers a fixed `avg_num_neighbors`. We can either compute
+            # this ~fixed statistic for OC20, or compute this value on-the-fly.
+            avg_num_neighbors=avg_num_neighbors,
+            rbf_hidden_channels=rbf_hidden_channels,
+            node_feats_down_irreps=o3.Irreps(node_feats_down_irreps),
+            edge_gates_irreps=edge_gates_irreps,
+            num_gates=num_gates,
+            multi_conv=multi_conv,
+            exponential=exponential,
+            use_source_target_edge_feats=edge_gates_use_source_target_feats,
+        )
+        self.interactions = torch.nn.ModuleList([inter])
+
+        # Use the appropriate self connection at the first layer for proper E0
+        use_sc_first = False
+        if "Residual" in str(interaction_cls_first):
+            use_sc_first = True
+
+        node_feats_irreps_out = inter.target_irreps
+        prod = EquivariantProductBasisBlock(
+            node_feats_irreps=node_feats_irreps_out,
+            target_irreps=hidden_irreps,
+            correlation=correlation,
+            num_elements=num_elements,
+            use_sc=use_sc_first,
+            contraction_type=contraction_type,
+        )
+        self.products = torch.nn.ModuleList([prod])
+
+        self.readouts = torch.nn.ModuleList()
+        self.readouts.append(LinearReadoutBlock(hidden_irreps))
+
+        for i in range(num_interactions - 1):
+            if i == num_interactions - 2:
+                if self.direct_forces:
+                    hidden_irreps_out = str(hidden_irreps[:2])
+                else:
+                    hidden_irreps_out = str(
+                        hidden_irreps[0]
+                    )  # Select only scalars for last layer
+            else:
+                hidden_irreps_out = hidden_irreps
+            inter = interaction_cls(
+                node_attrs_irreps=node_attr_irreps,
+                node_feats_irreps=hidden_irreps,
+                edge_attrs_irreps=sh_irreps,
+                edge_feats_irreps=edge_feats_irreps,
+                target_irreps=interaction_irreps,
+                hidden_irreps=hidden_irreps_out,
+                avg_num_neighbors=avg_num_neighbors,
+                rbf_hidden_channels=rbf_hidden_channels,
+                node_feats_down_irreps=o3.Irreps(node_feats_down_irreps),
+                edge_gates_irreps=edge_gates_irreps,
+                num_gates=num_gates,
+                multi_conv=multi_conv,
+                exponential=exponential,
+                use_source_target_edge_feats=edge_gates_use_source_target_feats,
+            )
+            self.interactions.append(inter)
+            prod = EquivariantProductBasisBlock(
+                node_feats_irreps=interaction_irreps,
+                target_irreps=hidden_irreps_out,
+                correlation=correlation,
+                num_elements=num_elements,
+                use_sc=True,
+                contraction_type=contraction_type,
+            )
+            self.products.append(prod)
+            if i == num_interactions - 2:
+                self.readouts.append(
+                    NonLinearReadoutBlock(hidden_irreps_out, MLP_irreps, gate)
+                )
+            else:
+                self.readouts.append(LinearReadoutBlock(hidden_irreps))
+
+    def atomic_numbers_to_compressed_one_hot(self, atomic_numbers):
+        """
+        Convert atomic numbers to compressed one-hot vectors.
+
+        Returns (num_atoms, num_elements),
+        where num_elements refers to the no. of elements present in the training
+        dataset, not an exhaustive list of all elements. For example, if the
+        only atomic numbers in the dataset are 1, 6, 8, 40, 80, then
+        num_elements would be 5 with the first column in the returned tensor
+        referring to hydrogen, 2nd to carbon and so on.
+        """
+        idx = self.z2idx[atomic_numbers]
+        assert torch.all(idx >= 0)
+
+        atomic_numbers_1hot = torch.zeros(
+            atomic_numbers.shape[0],
+            len(self.atomic_energies_fn.atomic_energies),
+            device=atomic_numbers.device,
+        ).scatter(1, idx.unsqueeze(1), 1.0)
+
+        return atomic_numbers_1hot
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data):
+
+        # TODO(@abhshkdz): Fit linear references per element from training data.
+        # These are currently initialized to 0.0.
+
+        # OCP prepro boilerplate.
+        pos = data.pos
+        atomic_numbers = data.atomic_numbers.long()
+        num_graphs = data.batch.max() + 1
+
+        # MACE computes forces via gradients.
+        pos.requires_grad_(True)
+
+        (
+            edge_index,
+            D_st,
+            distance_vec,
+            _,
+            _,
+            _,
+        ) = self.generate_graph(data)
+        ### OCP prepro ends.
+
+        # Atomic energies
+        #
+        # Comment(@abhshkdz): `data.node_attrs` is a 1-hot vector for each
+        # atomic number. `self.atomic_energies_fn` just matmuls the 1-hot
+        # vectors with the list of energies per atomic number, returning the
+        # energy per element.
+        atomic_numbers_1hot = self.atomic_numbers_to_compressed_one_hot(
+            atomic_numbers
+        )
+
+        node_e0 = self.atomic_energies_fn(atomic_numbers_1hot)
+        e0 = scatter_sum(
+            src=node_e0, index=data.batch, dim=-1, dim_size=num_graphs
+        )  # [n_graphs,]
+
+        # Embeddings
+        node_feats = self.node_embedding(atomic_numbers_1hot)
+
+        # Comment(@abhshkdz): `lengths` here is same as `D_st`, and `vectors` is
+        # the same as `distance_vec` but pointing in the opposite direction.
+        # vectors, lengths = get_edge_vectors_and_lengths(
+        #     positions=data.positions, edge_index=data.edge_index, shifts=data.shifts
+        # )
+        lengths = D_st.view(-1, 1)
+        vectors = -distance_vec
+
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(lengths)
+
+        # Interactions
+        energies = [e0]
+        for interaction, product, readout in zip(
+            self.interactions, self.products, self.readouts
+        ):
+            # print((neighbors / data.natoms).mean(), self.avg_num_neighbors)
+            node_feats, sc = interaction(
+                node_attrs=atomic_numbers_1hot,
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=edge_index,
+            )
+            node_feats = product(
+                node_feats=node_feats, sc=sc, node_attrs=atomic_numbers_1hot
+            )
+            node_energies = readout(node_feats).squeeze(-1)  # [n_nodes, ]
+            energy = scatter_sum(
+                src=node_energies,
+                index=data.batch,
+                dim=-1,
+                dim_size=num_graphs,
+            )  # [n_graphs,]
+            energies.append(energy)
+
+        # Sum over energy contributions
+        contributions = torch.stack(energies, dim=-1)
+        total_energy = torch.sum(contributions, dim=-1)  # [n_graphs, ]
+
+        # Compute forces via autograd.
+        forces = compute_forces(
+            energy=total_energy, positions=pos, training=self.training
+        )
+
+        return total_energy, forces
+
+
+@registry.register_model("scaleshift_mace")
+class ScaleShiftMACE(MACE):
+    def __init__(
+        self,
+        # Unused legacy OCP params.
+        num_atoms: Optional[int],
+        bond_feat_dim: int,
+        num_targets: int,
+        #
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        max_ell: int,
+        num_interactions: int,
+        num_elements: int,
+        hidden_irreps: str,
+        MLP_irreps: str,
+        avg_num_neighbors: float,
+        correlation: int,
+        atomic_inter_scale: float,
+        atomic_inter_shift: float,
+        # Defaults from OCP / https://github.com/ACEsuit/mace/blob/main/scripts/run_train.py
+        gate=torch.nn.functional.silu,
+        atomic_energies=str,
+        interaction_cls=RealAgnosticResidualInteractionBlock,
+        interaction_cls_first=RealAgnosticResidualInteractionBlock,
+        max_neighbors: int = 500,
+        otf_graph: bool = True,
+        use_pbc: bool = True,
+        regress_forces: bool = True,
+        # support for gaussian basis.
+        rbf: str = "bessel",
+        num_rbf: int = 8,
+        rbf_hidden_channels: int = 64,
+        # support for direct forces.
+        direct_forces: bool = False,
+    ):
+        super().__init__(
+            num_atoms,
+            bond_feat_dim,
+            num_targets,
+            #
+            r_max,
+            num_bessel,
+            num_polynomial_cutoff,
+            max_ell,
+            num_interactions,
+            num_elements,
+            hidden_irreps,
+            MLP_irreps,
+            avg_num_neighbors,
+            correlation,
+            gate,
+            atomic_energies,
+            interaction_cls,
+            interaction_cls_first,
+            max_neighbors,
+            otf_graph,
+            use_pbc,
+            regress_forces,
+            rbf=rbf,
+            num_rbf=num_rbf,
+            rbf_hidden_channels=rbf_hidden_channels,
+            direct_forces=direct_forces,
+        )
+        self.scale_shift = ScaleShiftBlock(
+            scale=atomic_inter_scale, shift=atomic_inter_shift
+        )
+
+        if self.direct_forces:
+            self.force_readout = ForceBlock(o3.Irreps(hidden_irreps))
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data):
+        # TODO(@abhshkdz): Fit linear references per element from training data.
+        # These are currently initialized to 0.0.
+
+        # OCP prepro boilerplate.
+        pos = data.pos
+        atomic_numbers = data.atomic_numbers.long()
+        num_graphs = data.batch.max() + 1
+
+        # MACE computes forces via gradients.
+        pos.requires_grad_(True)
+
+        (
+            edge_index,
+            D_st,
+            distance_vec,
+            _,
+            _,
+            _,
+        ) = self.generate_graph(data)
+
+        vectors = -distance_vec
+        lengths = D_st.view(-1, 1)
+        ### OCP prepro ends.
+
+        # Atomic energies
+        #
+        # Comment(@abhshkdz): `data.node_attrs` is a 1-hot vector for each
+        # atomic number. `self.atomic_energies_fn` just matmuls the 1-hot
+        # vectors with the list of energies per atomic number, returning the
+        # energy per element.
+        atomic_numbers_1hot = self.atomic_numbers_to_compressed_one_hot(
+            atomic_numbers
+        )
+
+        node_e0 = self.atomic_energies_fn(atomic_numbers_1hot)
+        e0 = scatter_sum(
+            src=node_e0, index=data.batch, dim=-1, dim_size=num_graphs
+        )  # [n_graphs,]
+
+        # Embeddings
+        node_feats = self.node_embedding(atomic_numbers_1hot)
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(lengths)
+
+        # Interactions
+        node_es_list = []
+        for interaction, product, readout in zip(
+            self.interactions, self.products, self.readouts
+        ):
+            node_feats, sc = interaction(
+                node_attrs=atomic_numbers_1hot,
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=edge_index,
+            )
+            node_feats = product(
+                node_feats=node_feats, sc=sc, node_attrs=atomic_numbers_1hot
+            )
+            node_es_list.append(
+                readout(node_feats).squeeze(-1)
+            )  # {[n_nodes, ], }
+
+        # Sum over interactions
+        node_inter_es = torch.sum(
+            torch.stack(node_es_list, dim=0), dim=0
+        )  # [n_nodes, ]
+        node_inter_es = self.scale_shift(node_inter_es)
+
+        # Sum over nodes in graph
+        inter_e = scatter_sum(
+            src=node_inter_es,
+            index=data.batch,
+            dim=-1,
+            dim_size=data.num_graphs,
+        )  # [n_graphs,]
+
+        # Add E_0 and (scaled) interaction energy
+        total_e = e0 + inter_e
+
+        if self.direct_forces:
+            forces = self.force_readout(node_feats)
+        else:
+            forces = compute_forces(
+                energy=inter_e, positions=pos, training=self.training
+            )
+
+        return total_e, forces
+
+
+@registry.register_model("interaction_energy_mace")
+class InteractionEnergyMACE(MACE):
+    def __init__(
+        self,
+        # Unused legacy OCP params.
+        num_atoms: Optional[int],
+        bond_feat_dim: int,
+        num_targets: int,
+        #
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        max_ell: int,
+        num_interactions: int,
+        num_elements: int,
+        hidden_irreps: str,
+        MLP_irreps: str,
+        avg_num_neighbors: float,
+        correlation: int,
+        # Defaults from OCP / https://github.com/ACEsuit/mace/blob/main/scripts/run_train.py
+        gate=torch.nn.functional.silu,
+        atomic_energies=str,
+        interaction_cls: str = "RealAgnosticResidualInteractionBlock",
+        interaction_cls_first: str = "RealAgnosticResidualInteractionBlock",
+        max_neighbors: int = 500,
+        otf_graph: bool = True,
+        use_pbc: bool = True,
+        regress_forces: bool = True,
+        # support for gaussian basis.
+        rbf: str = "bessel",
+        num_rbf: int = 8,
+        rbf_hidden_channels: int = 64,
+        # support for direct forces.
+        direct_forces: bool = False,
+        # support for species-agnostic contraction.
+        contraction_type: str = "v1",
+        # source and target feature size when concatenating for edge conv.
+        node_feats_down_irreps: str = "64x0e",
+        # parameters for edge gating.
+        edge_gates_irreps: str = "0e",
+        num_gates: int = 4,
+        multi_conv: bool = False,
+        exponential: bool = True,
+        edge_gates_use_source_target_feats: bool = False,
+    ):
+        super().__init__(
+            num_atoms,
+            bond_feat_dim,
+            num_targets,
+            #
+            r_max,
+            num_bessel,
+            num_polynomial_cutoff,
+            max_ell,
+            num_interactions,
+            num_elements,
+            hidden_irreps,
+            MLP_irreps,
+            avg_num_neighbors,
+            correlation,
+            gate,
+            atomic_energies,
+            eval(interaction_cls),
+            eval(interaction_cls_first),
+            max_neighbors,
+            otf_graph,
+            use_pbc,
+            regress_forces,
+            rbf=rbf,
+            num_rbf=num_rbf,
+            rbf_hidden_channels=rbf_hidden_channels,
+            direct_forces=direct_forces,
+            contraction_type=contraction_type,
+            node_feats_down_irreps=node_feats_down_irreps,
+            edge_gates_irreps=edge_gates_irreps,
+            num_gates=num_gates,
+            multi_conv=multi_conv,
+            exponential=exponential,
+            edge_gates_use_source_target_feats=edge_gates_use_source_target_feats,
+        )
+        if self.direct_forces:
+            self.force_readout = ForceBlock(o3.Irreps(hidden_irreps))
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data):
+        # TODO(@abhshkdz): Fit linear references per element from training data.
+        # These are currently initialized to 0.0.
+
+        # OCP prepro boilerplate.
+        pos = data.pos
+        atomic_numbers = data.atomic_numbers.long()
+
+        # MACE computes forces via gradients.
+        pos.requires_grad_(True)
+
+        (
+            edge_index,
+            D_st,
+            distance_vec,
+            _,
+            _,
+            _,
+        ) = self.generate_graph(data)
+
+        vectors = -distance_vec
+        lengths = D_st.view(-1, 1)
+        ### OCP prepro ends.
+
+        # Comment(@abhshkdz): `data.node_attrs` is a 1-hot vector for each
+        # atomic number. `self.atomic_energies_fn` just matmuls the 1-hot
+        # vectors with the list of energies per atomic number, returning the
+        # energy per element.
+        atomic_numbers_1hot = self.atomic_numbers_to_compressed_one_hot(
+            atomic_numbers
+        )
+
+        # Embeddings
+        node_feats = self.node_embedding(atomic_numbers_1hot)
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(lengths)
+
+        # Interactions
+        node_es_list = []
+        for interaction, product, readout in zip(
+            self.interactions, self.products, self.readouts
+        ):
+            node_feats, sc = interaction(
+                node_attrs=atomic_numbers_1hot,
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=edge_index,
+            )
+            node_feats = product(
+                node_feats=node_feats, sc=sc, node_attrs=atomic_numbers_1hot
+            )
+            node_es_list.append(
+                readout(node_feats).squeeze(-1)
+            )  # {[n_nodes, ], }
+
+        # Sum over interactions
+        node_inter_es = torch.sum(
+            torch.stack(node_es_list, dim=0), dim=0
+        )  # [n_nodes, ]
+
+        # Sum over nodes in graph
+        inter_e = scatter_sum(
+            src=node_inter_es,
+            index=data.batch,
+            dim=-1,
+            dim_size=data.num_graphs,
+        )  # [n_graphs,]
+
+        energy = inter_e
+
+        if self.direct_forces:
+            forces = self.force_readout(node_feats)
+            return energy, forces
+        else:
+            raise NotImplementedError
diff --git a/ocpmodels/models/mace/radial.py b/ocpmodels/models/mace/radial.py
new file mode 100644
index 0000000000..59bd477b50
--- /dev/null
+++ b/ocpmodels/models/mace/radial.py
@@ -0,0 +1,113 @@
+import numpy as np
+import torch
+
+
+class BesselBasis(torch.nn.Module):
+    """
+    Klicpera, J.; Groß, J.; Günnemann, S. Directional Message Passing for Molecular Graphs; ICLR 2020.
+    Equation (7)
+    """
+
+    def __init__(self, r_max: float, num_basis=8, trainable=False):
+        super().__init__()
+
+        bessel_weights = (
+            np.pi
+            / r_max
+            * torch.linspace(
+                start=1.0,
+                end=num_basis,
+                steps=num_basis,
+                dtype=torch.get_default_dtype(),
+            )
+        )
+        if trainable:
+            self.bessel_weights = torch.nn.Parameter(bessel_weights)
+        else:
+            self.register_buffer("bessel_weights", bessel_weights)
+
+        self.register_buffer(
+            "r_max", torch.tensor(r_max, dtype=torch.get_default_dtype())
+        )
+        self.register_buffer(
+            "prefactor",
+            torch.tensor(
+                np.sqrt(2.0 / r_max), dtype=torch.get_default_dtype()
+            ),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:  # [..., 1]
+        numerator = torch.sin(self.bessel_weights * x)  # [..., num_basis]
+        return self.prefactor * (numerator / x)
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(r_max={self.r_max}, num_basis={len(self.bessel_weights)}, "
+            f"trainable={self.bessel_weights.requires_grad})"
+        )
+
+
+# From https://github.com/Open-Catalyst-Project/ocp/blob/a8ac36d75fc60d5b6d54b4019502941eb868fed0/ocpmodels/models/gemnet_oc/layers/radial_basis.py#L65
+class GaussianBasis(torch.nn.Module):
+    def __init__(self, r_max: float, num_gaussians=50, trainable=False):
+        super().__init__()
+
+        offset = torch.linspace(0, r_max, num_gaussians)
+        if trainable:
+            self.offset = torch.nn.Parameter(offset, requires_grad=True)
+        else:
+            self.register_buffer("offset", offset)
+
+        self.register_buffer(
+            "coeff",
+            torch.tensor(
+                -0.5 / (r_max / (num_gaussians - 1)) ** 2,
+                dtype=torch.get_default_dtype(),
+            ),
+        )
+
+        self.register_buffer(
+            "r_max", torch.tensor(r_max, dtype=torch.get_default_dtype())
+        )
+
+    def forward(self, dist: torch.Tensor) -> torch.Tensor:
+        dist = dist - self.offset[None, :]
+        return torch.exp(self.coeff * torch.pow(dist, 2))
+
+
+class PolynomialCutoff(torch.nn.Module):
+    """
+    Klicpera, J.; Groß, J.; Günnemann, S. Directional Message Passing for Molecular Graphs; ICLR 2020.
+    Equation (8)
+    """
+
+    p: torch.Tensor
+    r_max: torch.Tensor
+
+    def __init__(self, r_max: float, p=6):
+        super().__init__()
+        self.register_buffer(
+            "p", torch.tensor(p, dtype=torch.get_default_dtype())
+        )
+        self.register_buffer(
+            "r_max", torch.tensor(r_max, dtype=torch.get_default_dtype())
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # yapf: disable
+        envelope = (
+                1.0
+                - ((self.p + 1.0) * (self.p + 2.0) / 2.0) * torch.pow(x / self.r_max, self.p)
+                + self.p * (self.p + 2.0) * torch.pow(x / self.r_max, self.p + 1)
+                - (self.p * (self.p + 1.0) / 2) * torch.pow(x / self.r_max, self.p + 2)
+        )
+        # yapf: enable
+
+        # noinspection PyUnresolvedReferences
+        return envelope * (x < self.r_max).type(torch.get_default_dtype())
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(p={self.p}, r_max={self.r_max})"
diff --git a/ocpmodels/models/mace/tools/torch_tools.py b/ocpmodels/models/mace/tools/torch_tools.py
new file mode 100644
index 0000000000..dc7f40777e
--- /dev/null
+++ b/ocpmodels/models/mace/tools/torch_tools.py
@@ -0,0 +1,72 @@
+import logging
+from typing import Dict
+
+import numpy as np
+import torch
+
+TensorDict = Dict[str, torch.Tensor]
+
+
+def to_one_hot(indices: torch.Tensor, num_classes: int) -> torch.Tensor:
+    """
+    Generates one-hot encoding with <num_classes> classes from <indices>
+    :param indices: (N x 1) tensor
+    :param num_classes: number of classes
+    :param device: torch device
+    :return: (N x num_classes) tensor
+    """
+    shape = indices.shape[:-1] + (num_classes,)
+    oh = torch.zeros(shape, device=indices.device).view(shape)
+
+    # scatter_ is the in-place version of scatter
+    oh.scatter_(dim=-1, index=indices, value=1)
+
+    return oh.view(*shape)
+
+
+def count_parameters(module: torch.nn.Module) -> int:
+    return int(sum(np.prod(p.shape) for p in module.parameters()))
+
+
+def tensor_dict_to_device(td: TensorDict, device: torch.device) -> TensorDict:
+    return {k: v.to(device) for k, v in td.items()}
+
+
+def set_seeds(seed: int) -> None:
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def to_numpy(t: torch.Tensor) -> np.ndarray:
+    return t.cpu().detach().numpy()
+
+
+def init_device(device_str: str) -> torch.device:
+    if device_str == "cuda":
+        assert torch.cuda.is_available(), "No CUDA device available!"
+        logging.info(
+            f"CUDA version: {torch.version.cuda}, CUDA device: {torch.cuda.current_device()}"
+        )
+        torch.cuda.init()
+        return torch.device("cuda")
+
+    logging.info("Using CPU")
+    return torch.device("cpu")
+
+
+dtype_dict = {"float32": torch.float32, "float64": torch.float64}
+
+
+def set_default_dtype(dtype: str) -> None:
+    torch.set_default_dtype(dtype_dict[dtype])
+
+
+def get_complex_default_dtype():
+    default_dtype = torch.get_default_dtype()
+    if default_dtype == torch.float64:
+        return torch.complex128
+
+    if default_dtype == torch.float32:
+        return torch.complex64
+
+    raise NotImplementedError
diff --git a/ocpmodels/models/mace/utils.py b/ocpmodels/models/mace/utils.py
new file mode 100644
index 0000000000..68c20f319d
--- /dev/null
+++ b/ocpmodels/models/mace/utils.py
@@ -0,0 +1,120 @@
+import logging
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn
+import torch.utils.data
+
+from .tools.torch_tools import to_numpy
+from .mace_core.scatter import scatter_sum
+
+from .blocks import AtomicEnergiesBlock
+
+
+def compute_forces(
+    energy: torch.Tensor, positions: torch.Tensor, training=True
+) -> torch.Tensor:
+    gradient = torch.autograd.grad(
+        outputs=energy,  # [n_graphs, ]
+        inputs=positions,  # [n_nodes, 3]
+        grad_outputs=torch.ones_like(energy),
+        retain_graph=training,  # Make sure the graph is not destroyed during training
+        create_graph=training,  # Create graph for second derivative
+        only_inputs=True,  # Diff only w.r.t. inputs
+        allow_unused=True,
+    )[
+        0
+    ]  # [n_nodes, 3]
+    if gradient is None:
+        logging.warning("Gradient is None, padded with zeros")
+        return torch.zeros_like(positions)
+    return -1 * gradient
+
+
+def get_edge_vectors_and_lengths(
+    positions: torch.Tensor,  # [n_nodes, 3]
+    edge_index: torch.Tensor,  # [2, n_edges]
+    shifts: torch.Tensor,  # [n_edges, 3]
+    normalize: bool = False,
+    eps: float = 1e-9,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    sender, receiver = edge_index
+    # From ase.neighborlist:
+    # D = positions[j]-positions[i]+S.dot(cell)
+    # where shifts = S.dot(cell)
+    vectors = positions[receiver] - positions[sender] + shifts  # [n_edges, 3]
+    lengths = torch.linalg.norm(vectors, dim=-1, keepdim=True)  # [n_edges, 1]
+    if normalize:
+        vectors_normed = vectors / (lengths + eps)
+        return vectors_normed, lengths
+
+    return vectors, lengths
+
+
+def compute_mean_std_atomic_inter_energy(
+    data_loader: torch.utils.data.DataLoader,
+    atomic_energies: np.ndarray,
+) -> Tuple[float, float]:
+    atomic_energies_fn = AtomicEnergiesBlock(atomic_energies=atomic_energies)
+
+    avg_atom_inter_es_list = []
+
+    for batch in data_loader:
+        node_e0 = atomic_energies_fn(batch.node_attrs)
+        graph_e0s = scatter_sum(
+            src=node_e0, index=batch.batch, dim=-1, dim_size=batch.num_graphs
+        )
+        graph_sizes = batch.ptr[1:] - batch.ptr[:-1]
+        avg_atom_inter_es_list.append(
+            (batch.energy - graph_e0s) / graph_sizes
+        )  # {[n_graphs], }
+
+    avg_atom_inter_es = torch.cat(avg_atom_inter_es_list)  # [total_n_graphs]
+    mean = to_numpy(torch.mean(avg_atom_inter_es)).item()
+    std = to_numpy(torch.std(avg_atom_inter_es)).item()
+
+    return mean, std
+
+
+def compute_mean_rms_energy_forces(
+    data_loader: torch.utils.data.DataLoader,
+    atomic_energies: np.ndarray,
+) -> Tuple[float, float]:
+    atomic_energies_fn = AtomicEnergiesBlock(atomic_energies=atomic_energies)
+
+    atom_energy_list = []
+    forces_list = []
+
+    for batch in data_loader:
+        node_e0 = atomic_energies_fn(batch.node_attrs)
+        graph_e0s = scatter_sum(
+            src=node_e0, index=batch.batch, dim=-1, dim_size=batch.num_graphs
+        )
+        graph_sizes = batch.ptr[1:] - batch.ptr[:-1]
+        atom_energy_list.append(
+            (batch.energy - graph_e0s) / graph_sizes
+        )  # {[n_graphs], }
+        forces_list.append(batch.forces)  # {[n_graphs*n_atoms,3], }
+
+    atom_energies = torch.cat(atom_energy_list, dim=0)  # [total_n_graphs]
+    forces = torch.cat(forces_list, dim=0)  # {[total_n_graphs*n_atoms,3], }
+
+    mean = to_numpy(torch.mean(atom_energies)).item()
+    rms = to_numpy(torch.sqrt(torch.mean(torch.square(forces)))).item()
+
+    return mean, rms
+
+
+def compute_avg_num_neighbors(data_loader: torch.utils.data.DataLoader) -> float:
+    num_neighbors = []
+
+    for batch in data_loader:
+        _, receivers = batch.edge_index
+        _, counts = torch.unique(receivers, return_counts=True)
+        num_neighbors.append(counts)
+
+    avg_num_neighbors = torch.mean(
+        torch.cat(num_neighbors, dim=0).type(torch.get_default_dtype())
+    )
+    return to_numpy(avg_num_neighbors).item()
diff --git a/ocpmodels/modules/evaluator.py b/ocpmodels/modules/evaluator.py
index 655f981694..4cc8c37eeb 100644
--- a/ocpmodels/modules/evaluator.py
+++ b/ocpmodels/modules/evaluator.py
@@ -38,7 +38,9 @@ class Evaluator:
             "forces_mae",
             "forces_cos",
             "forces_magnitude",
+            "forces_mse",
             "energy_mae",
+            "energy_per_atom_mse",
             "energy_force_within_threshold",
         ],
         "is2rs": [
@@ -116,6 +118,13 @@ def energy_mse(prediction, target):
     return squared_error(prediction["energy"], target["energy"])
 
 
+def energy_per_atom_mse(prediction, target):
+    return squared_error(
+        prediction["energy"] / prediction["natoms"],
+        target["energy"] / target["natoms"]
+    )
+
+
 def forcesx_mae(prediction, target):
     return absolute_error(prediction["forces"][:, 0], target["forces"][:, 0])
 
diff --git a/ocpmodels/modules/loss.py b/ocpmodels/modules/loss.py
index 184e6bb161..abe746c16d 100644
--- a/ocpmodels/modules/loss.py
+++ b/ocpmodels/modules/loss.py
@@ -6,26 +6,54 @@
 from ocpmodels.common import distutils
 
 
-class L2MAELoss(nn.Module):
+class OCPLoss(nn.Module):
     def __init__(self, reduction="mean"):
         super().__init__()
         self.reduction = reduction
         assert reduction in ["mean", "sum"]
 
     def forward(self, input: torch.Tensor, target: torch.Tensor):
-        dists = torch.norm(input - target, p=2, dim=-1)
+        raise NotImplementedError
+
+
+# Following `weighted_mean_squared_error_energy` from MACE
+# https://github.com/ACEsuit/mace/blob/main/mace/modules/loss.py#L18
+class PerAtomMSELossEnergy(OCPLoss):
+    def forward(
+        self,
+        input: torch.Tensor,
+        target: torch.Tensor,
+        natoms: torch.Tensor,
+    ):
+        dists = torch.square((target - input) / natoms)  # (nGraphs, )
         if self.reduction == "mean":
             return torch.mean(dists)
         elif self.reduction == "sum":
             return torch.sum(dists)
 
 
-class AtomwiseL2Loss(nn.Module):
-    def __init__(self, reduction="mean"):
-        super().__init__()
-        self.reduction = reduction
-        assert reduction in ["mean", "sum"]
+# Following `mean_squared_error_forces` from MACE
+# https://github.com/ACEsuit/mace/blob/main/mace/modules/loss.py#L54
+class MSELossForces(OCPLoss):
+    def forward(self, input: torch.Tensor, target: torch.Tensor):
+        dists = torch.square(target - input)  # (nAtoms, 3)
+        dists = torch.mean(dists, dim=-1)  # (nAtoms, )
+        if self.reduction == "mean":
+            return torch.mean(dists)
+        elif self.reduction == "sum":
+            return torch.sum(dists)
+
+
+class L2MAELoss(OCPLoss):
+    def forward(self, input: torch.Tensor, target: torch.Tensor):
+        dists = torch.norm(input - target, p=2, dim=-1)
+        if self.reduction == "mean":
+            return torch.mean(dists)
+        elif self.reduction == "sum":
+            return torch.sum(dists)
+
 
+class AtomwiseL2Loss(OCPLoss):
     def forward(
         self,
         input: torch.Tensor,
diff --git a/ocpmodels/trainers/eval_mace_chkpts_trainer.py b/ocpmodels/trainers/eval_mace_chkpts_trainer.py
new file mode 100644
index 0000000000..4c2f9a8ab4
--- /dev/null
+++ b/ocpmodels/trainers/eval_mace_chkpts_trainer.py
@@ -0,0 +1,60 @@
+import os
+import errno
+import logging
+
+import torch
+import torch.optim as optim
+
+from ocpmodels.common.registry import registry
+from ocpmodels.common.utils import load_state_dict
+from ocpmodels.trainers.mace_trainer import MACETrainer
+
+
+@registry.register_trainer("eval_mace_chkpts_trainer")
+class EvalMACECheckpointsTrainer(MACETrainer):
+    def load_checkpoint(self, checkpoint_path):
+        if not os.path.isfile(checkpoint_path):
+            raise FileNotFoundError(
+                errno.ENOENT, "Checkpoint file not found", checkpoint_path
+            )
+
+        logging.info(f"Loading checkpoint from: {checkpoint_path}")
+        map_location = torch.device("cpu") if self.cpu else self.device
+        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+        self.epoch = checkpoint.get("epoch", 0)
+        self.step = checkpoint.get("step", 0)
+        self.best_val_metric = checkpoint.get("best_val_metric", None)
+        self.primary_metric = checkpoint.get("primary_metric", None)
+
+        # Match the "module." count in the keys of model and checkpoint state_dict
+        # DataParallel model has 1 "module.",  DistributedDataParallel has 2 "module."
+        # Not using either of the above two would have no "module."
+
+        ckpt_key_count = next(iter(checkpoint["model"])).count("module")
+        mod_key_count = next(iter(self.model.state_dict())).count("module")
+        key_count_diff = mod_key_count - ckpt_key_count
+
+        if key_count_diff > 0:
+            new_dict = {
+                key_count_diff * "module." + k: v
+                for k, v in checkpoint["model"].items()
+            }
+        elif key_count_diff < 0:
+            new_dict = {
+                k[len("module.") * abs(key_count_diff) :]: v
+                for k, v in checkpoint["model"].items()
+            }
+        else:
+            new_dict = checkpoint["model"]
+
+        strict = self.config["task"].get("strict_load", True)
+        load_state_dict(self.model, new_dict, strict=strict)
+
+        if "optimizer" in checkpoint:
+            self.optimizer.load_state_dict(checkpoint["optimizer"])
+        if "lr_scheduler" in checkpoint and checkpoint["lr_scheduler"] is not None:
+            self.scheduler.scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        if "ema" in checkpoint and checkpoint["ema"] is not None:
+            self.ema.load_state_dict(checkpoint["ema"])
+        else:
+            self.ema = None
diff --git a/ocpmodels/trainers/mace_trainer.py b/ocpmodels/trainers/mace_trainer.py
new file mode 100644
index 0000000000..02a1ae43f1
--- /dev/null
+++ b/ocpmodels/trainers/mace_trainer.py
@@ -0,0 +1,265 @@
+import logging
+
+import torch
+import torch.optim as optim
+
+from ocpmodels.common import distutils
+from ocpmodels.common.registry import registry
+from ocpmodels.modules.loss import (
+    DDPLoss,
+    L2MAELoss,
+    MSELossForces,
+    PerAtomMSELossEnergy,
+)
+from ocpmodels.trainers import ForcesTrainer
+
+
+@registry.register_trainer("mace_trainer")
+class MACETrainer(ForcesTrainer):
+    # Following optimizer / weight decay defaults from
+    # https://github.com/ACEsuit/mace/blob/5470b632d839358faed4e9c97f67fece1b558962/scripts/run_train.py#L210
+    def load_optimizer(self):
+        optimizer = self.config["optim"].get("optimizer", "AdamW")
+        optimizer = getattr(optim, optimizer)
+
+        if self.config["optim"].get("weight_decay", 0) > 0:
+
+            int_params_decay = []
+            int_params_no_decay = []
+            for (
+                name,
+                param,
+            ) in self._unwrapped_model.interactions.named_parameters():
+                if "linear.weight" in name or "skip_tp_full.weight" in name:
+                    int_params_decay += [param]
+                else:
+                    int_params_no_decay += [param]
+
+            self.optimizer = optimizer(
+                params=[
+                    {
+                        "name": "embedding",
+                        "params": self._unwrapped_model.node_embedding.parameters(),
+                        "weight_decay": 0.0,
+                    },
+                    {
+                        "name": "interactions_decay",
+                        "params": int_params_decay,
+                        "weight_decay": self.config["optim"]["weight_decay"],
+                    },
+                    {
+                        "name": "interactions_no_decay",
+                        "params": int_params_no_decay,
+                        "weight_decay": 0.0,
+                    },
+                    {
+                        "name": "products",
+                        "params": self._unwrapped_model.products.parameters(),
+                        "weight_decay": self.config["optim"]["weight_decay"],
+                    },
+                    {
+                        "name": "readouts",
+                        "params": self._unwrapped_model.readouts.parameters(),
+                        "weight_decay": 0.0,
+                    },
+                ],
+                lr=self.config["optim"]["lr_initial"],
+                **self.config["optim"].get("optimizer_params", {}),
+            )
+        else:
+            self.optimizer = optimizer(
+                params=self.model.parameters(),
+                lr=self.config["optim"]["lr_initial"],
+                **self.config["optim"].get("optimizer_params", {}),
+            )
+
+    def load_loss(self):
+        self.loss_fn = {}
+        self.loss_fn["energy"] = DDPLoss(PerAtomMSELossEnergy())
+        self.loss_fn["force"] = DDPLoss(MSELossForces())
+
+    def _compute_loss(self, out, batch_list):
+        assert self.config["model_attributes"].get("regress_forces", False)
+
+        loss = []
+
+        natoms = torch.cat(
+            [batch.natoms.to(self.device) for batch in batch_list]
+        )
+
+        # Energy loss.
+        energy_target = torch.cat(
+            [batch.y.to(self.device) for batch in batch_list], dim=0
+        )
+        if self.normalizer.get("normalize_labels", False):
+            energy_target = self.normalizers["target"].norm(energy_target)
+        energy_mult = self.config["optim"].get("energy_coefficient", 1)
+        loss.append(
+            energy_mult
+            * self.loss_fn["energy"](out["energy"], energy_target, natoms)
+        )
+
+        # Force loss.
+        force_target = torch.cat(
+            [batch.force.to(self.device) for batch in batch_list], dim=0
+        )
+        if self.normalizer.get("normalize_labels", False):
+            force_target = self.normalizers["grad_target"].norm(force_target)
+
+        force_mult = self.config["optim"].get("force_coefficient", 100)
+
+        if self.config["task"].get("train_on_free_atoms", False):
+            fixed = torch.cat(
+                [batch.fixed.to(self.device) for batch in batch_list]
+            )
+            mask = fixed == 0
+            loss.append(
+                force_mult
+                * self.loss_fn["force"](
+                    out["forces"][mask], force_target[mask]
+                )
+            )
+        else:
+            loss.append(
+                force_mult * self.loss_fn["force"](out["forces"], force_target)
+            )
+
+        # Sanity check to make sure the compute graph is correct.
+        for lc in loss:
+            assert hasattr(lc, "grad_fn")
+
+        loss = sum(loss)
+        return loss
+
+
+@registry.register_trainer("interaction_mace_trainer")
+class InteractionMACETrainer(MACETrainer):
+    def load_loss(self):
+        self.loss_fn = {}
+        self.loss_fn["energy"] = DDPLoss(torch.nn.L1Loss())
+        self.loss_fn["force"] = DDPLoss(L2MAELoss())
+
+    def _compute_loss(self, out, batch_list):
+        loss = []
+
+        # Energy loss.
+        energy_target = torch.cat(
+            [batch.y.to(self.device) for batch in batch_list], dim=0
+        )
+        if self.normalizer.get("normalize_labels", False):
+            energy_target = self.normalizers["target"].norm(energy_target)
+        energy_mult = self.config["optim"].get("energy_coefficient", 1)
+        loss.append(
+            energy_mult * self.loss_fn["energy"](out["energy"], energy_target)
+        )
+
+        # Force loss.
+        if self.config["model_attributes"].get("regress_forces", True):
+            force_target = torch.cat(
+                [batch.force.to(self.device) for batch in batch_list], dim=0
+            )
+            if self.normalizer.get("normalize_labels", False):
+                force_target = self.normalizers["grad_target"].norm(
+                    force_target
+                )
+
+            tag_specific_weights = self.config["task"].get(
+                "tag_specific_weights", []
+            )
+            if tag_specific_weights != []:
+                # handle tag specific weights as introduced in forcenet
+                assert len(tag_specific_weights) == 3
+
+                batch_tags = torch.cat(
+                    [
+                        batch.tags.float().to(self.device)
+                        for batch in batch_list
+                    ],
+                    dim=0,
+                )
+                weight = torch.zeros_like(batch_tags)
+                weight[batch_tags == 0] = tag_specific_weights[0]
+                weight[batch_tags == 1] = tag_specific_weights[1]
+                weight[batch_tags == 2] = tag_specific_weights[2]
+
+                if self.config["optim"].get("loss_force", "l2mae") == "l2mae":
+                    # zero out nans, if any
+                    found_nans_or_infs = not torch.all(
+                        out["forces"].isfinite()
+                    )
+                    if found_nans_or_infs is True:
+                        logging.warning("Found nans while computing loss")
+                        out["forces"] = torch.nan_to_num(
+                            out["forces"], nan=0.0
+                        )
+
+                    dists = torch.norm(
+                        out["forces"] - force_target, p=2, dim=-1
+                    )
+                    weighted_dists_sum = (dists * weight).sum()
+
+                    num_samples = out["forces"].shape[0]
+                    num_samples = distutils.all_reduce(
+                        num_samples, device=self.device
+                    )
+                    weighted_dists_sum = (
+                        weighted_dists_sum
+                        * distutils.get_world_size()
+                        / num_samples
+                    )
+
+                    force_mult = self.config["optim"].get(
+                        "force_coefficient", 30
+                    )
+                    loss.append(force_mult * weighted_dists_sum)
+                else:
+                    raise NotImplementedError
+            else:
+                # Force coefficient = 30 has been working well for us.
+                force_mult = self.config["optim"].get("force_coefficient", 30)
+                if self.config["task"].get("train_on_free_atoms", False):
+                    fixed = torch.cat(
+                        [batch.fixed.to(self.device) for batch in batch_list]
+                    )
+                    mask = fixed == 0
+                    if (
+                        self.config["optim"]
+                        .get("loss_force", "mae")
+                        .startswith("atomwise")
+                    ):
+                        force_mult = self.config["optim"].get(
+                            "force_coefficient", 1
+                        )
+                        natoms = torch.cat(
+                            [
+                                batch.natoms.to(self.device)
+                                for batch in batch_list
+                            ]
+                        )
+                        natoms = torch.repeat_interleave(natoms, natoms)
+                        force_loss = force_mult * self.loss_fn["force"](
+                            out["forces"][mask],
+                            force_target[mask],
+                            natoms=natoms[mask],
+                            batch_size=batch_list[0].natoms.shape[0],
+                        )
+                        loss.append(force_loss)
+                    else:
+                        loss.append(
+                            force_mult
+                            * self.loss_fn["force"](
+                                out["forces"][mask], force_target[mask]
+                            )
+                        )
+                else:
+                    loss.append(
+                        force_mult
+                        * self.loss_fn["force"](out["forces"], force_target)
+                    )
+
+        # Sanity check to make sure the compute graph is correct.
+        for lc in loss:
+            assert hasattr(lc, "grad_fn")
+
+        loss = sum(loss)
+        return loss
diff --git a/scripts/compute_mace_stats.py b/scripts/compute_mace_stats.py
new file mode 100644
index 0000000000..b08c1ce6a7
--- /dev/null
+++ b/scripts/compute_mace_stats.py
@@ -0,0 +1,94 @@
+from ocpmodels.datasets import OC22LmdbDataset, data_list_collater
+from ocpmodels.models import BaseModel
+from torch.utils.data import DataLoader
+from torch_scatter import scatter
+from tqdm import tqdm
+import numpy as np
+import torch
+
+
+dset_path = "/checkpoint/bmwood/oc20_small/oc20_xxs_dataset/lmdb/train"
+oc20_ref = "/checkpoint/janlan/ocp/other_data/final_ref_energies_02_07_2021.pkl"
+cutoff = 6.0
+max_neighbors = 500
+
+# adapted from https://github.com/ACEsuit/mace/blob/f304c07bbafd651f5a52a447c954f1ef561d42e2/mace/data/utils.py#L170
+def compute_average_E0s(dl: DataLoader):
+    len_dset = len(dl.dataset)
+    len_zs = 83
+
+    A = np.zeros((len_dset, len_zs))
+    B = np.zeros(len_dset)
+
+    for i, batch in tqdm(enumerate(dl)):
+        breakpoint()
+        if i == len(dl) - 1: sz = batch.y.shape[0]
+        else: sz = dl.batch_size
+
+        B[i * dl.batch_size : i * dl.batch_size + sz] = batch.y
+        for j, do in enumerate(batch.to_data_list()):
+            for z in range(len_zs):
+                A[i * dl.batch_size + j, z] = np.count_nonzero(do.atomic_numbers - 1 == z)  # -1 because of 0-indexing
+
+
+    E0s = np.linalg.lstsq(A, B, rcond=None)[0]
+    return E0s
+
+
+if __name__ == "__main__":
+    cfg = {
+        "src": dset_path,
+        "oc20_ref": oc20_ref,
+        "total_energy": True,
+        "otf_graph": True,
+    }
+    dset = OC22LmdbDataset(cfg)
+
+    dl = DataLoader(
+        dset,
+        batch_size=32,
+        shuffle=False,
+        collate_fn=data_list_collater,
+        num_workers=32,
+        pin_memory=True,
+    )
+
+    E0s = compute_average_E0s(dl)
+    print("E0s", E0s.tolist())
+
+    E0s = torch.tensor(E0s)
+
+    model = BaseModel()
+
+    e, f, n, nbrs = [], [], [], []
+
+    for i, batch in tqdm(enumerate(dl)):
+
+        n.append(batch.natoms)
+        f.append(batch.force)
+
+        graph = model.generate_graph(
+            batch,
+            cutoff=cutoff,
+            max_neighbors=max_neighbors,
+            use_pbc=True,
+            otf_graph=True,
+        )
+
+        nbrs.append(torch.div(graph[-1], batch.natoms))
+
+        # compute interaction energies
+        atom_e0 = E0s[batch.atomic_numbers.long() - 1]
+        graph_e0 = scatter(src=atom_e0, index=batch.batch, dim=-1)
+
+        e.append(batch.y - graph_e0)
+
+    e = torch.cat(e, dim=0)
+    f = torch.cat(f, dim=0)
+    n = torch.cat(n, dim=0)
+    nbrs = torch.cat(nbrs, dim=0)
+
+    print("interaction energy", e.mean().item())
+    print("interaction energy per atom", torch.div(e, n).mean().item())
+    print("forces rms", torch.sqrt(torch.mean(torch.square(f))).item())
+    print("neighbors", nbrs.mean().item())
diff --git a/scripts/run_train.sh b/scripts/run_train.sh
new file mode 100755
index 0000000000..c147fa0cd6
--- /dev/null
+++ b/scripts/run_train.sh
@@ -0,0 +1,86 @@
+#!/bin/sh
+PATH=/checkpoint/abhshkdz/local/bin:/private/home/abhshkdz/.local/bin:/public/apps/anaconda3/2020.11/condabin:/public/apps/anaconda3/2020.11/bin:/public/apps/gcc/7.1.0/bin:/public/apps/cuda/11.6/bin:/usr/local/cuda/bin:/opt/bin:/usr/local/cuda/bin:/opt/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/public/slurm/20.11.5/bin:/public/slurm/20.11.5/sbin
+
+################################################################################
+
+ROOT_DIR="/private/home/abhshkdz/projects/mace-ocp"
+
+# name your run; it'll show up on wandb by this name
+IDENTIFIER="mace_0127"
+
+# yaml config
+CONFIG_YML=configs/s2ef/2M/mace/base_0127.yml
+
+# directory where slurm logs will get saved
+LOG_DIR=/checkpoint/abhshkdz/ocp_oct1_logs
+
+# conda env
+ENV=/private/home/abhshkdz/.conda/envs/ocp-mace-2022dec20
+
+# running locally?
+LOCAL_RUN=false
+
+# slurm?
+PARTITION=ocp,learnaccel
+NODES=2
+
+# initializing from a checkpoint?
+# CHECKPOINT_PATH=
+
+MODE=train
+
+# launching a sweep?
+# SWEEP_YML=experimental/abhshkdz/gnoc_v2/configs/2M/101103_no_loss_high_fmax_sweep.yml
+
+################################################################################
+
+PYTHON=$ENV/bin/python
+
+_setArgs(){
+    while [ "$1" != "" ]; do
+        case $1 in
+            "--local")
+                IDENTIFIER="debug"
+                LOCAL_RUN=true
+                ;;
+        esac
+        shift
+    done
+}
+
+_setArgs $*
+
+################################################################################
+
+LAUNCH_ARGS="--config-yml $CONFIG_YML \
+    --identifier $IDENTIFIER \
+    --logdir $LOG_DIR \
+    --seed 1 \
+    --distributed \
+    --mode $MODE"
+
+# LAUNCH_ARGS="$LAUNCH_ARGS --amp"
+
+# with pretrained checkpoint
+if [ ${CHECKPOINT_PATH+x} ] ; then
+    LAUNCH_ARGS="$LAUNCH_ARGS --checkpoint $CHECKPOINT_PATH"
+fi
+
+# with sweep params
+if [ ${SWEEP_YML+x} ] ; then
+    LAUNCH_ARGS="$LAUNCH_ARGS --sweep-yml $SWEEP_YML"
+fi
+
+if [ "$LOCAL_RUN" = true ] ; then
+    # local run
+    LAUNCH_ARGS="$LAUNCH_ARGS --debug"
+    $PYTHON -u -m torch.distributed.launch --nproc_per_node=1 main.py $LAUNCH_ARGS
+else
+    # submit to slurm
+    LAUNCH_ARGS="$LAUNCH_ARGS --submit \
+        --slurm-mem 480 \
+        --slurm-partition $PARTITION \
+        --num-nodes $NODES \
+        --num-gpus 8"
+    $PYTHON $ROOT_DIR/main.py $LAUNCH_ARGS
+fi