Error in OCPCalculator when using most checkpoints from OC22

[gauthierlab]

Hi there -- when using a checkpoint from OC22 (all of them except for GemNet-dT on this page: https://github.com/Open-Catalyst-Project/ocp/blob/main/MODELS.md#open-catalyst-2022-oc22) I get the following traceback in OCPCalculator:

---

WARNING:root:Unable to identify OCP trainer, defaulting to forces. Specify the trainer argument into OCPCalculator if otherwise.
Unknown option: -C
usage: git [--version] [--help] [-c name=value]
[--exec-path[=]] [--html-path] [--man-path] [--info-path]
[-p|--paginate|--no-pager] [--no-replace-objects] [--bare]
[--git-dir=] [--work-tree=] [--namespace=]
[]
WARNING:root:Unrecognized arguments: ['symmetric_edge_symmetrization']
Traceback (most recent call last):
File "/home/josegaut/n2_oxidation/test_ocp/breaking_oc22/./mlrelax.py", line 24, in
opt.run(fmax=0.05, steps=100)
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/ase/optimize/optimize.py", line 269, in run
return Dynamics.run(self)
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/ase/optimize/optimize.py", line 156, in run
for converged in Dynamics.irun(self):
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/ase/optimize/optimize.py", line 122, in irun
self.atoms.get_forces()
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/ase/atoms.py", line 788, in get_forces
forces = self._calc.get_forces(self)
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/ase/calculators/abc.py", line 23, in get_forces
return self.get_property('forces', atoms)
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/ase/calculators/calculator.py", line 737, in get_property
self.calculate(atoms, [name], system_changes)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/common/relaxation/ase_utils.py", line 205, in calculate
predictions = self.trainer.predict(
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/trainers/forces_trainer.py", line 195, in predict
out = self._forward(batch_list)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/trainers/forces_trainer.py", line 446, in _forward
out_energy, out_forces = self.model(batch_list)
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/common/data_parallel.py", line 58, in forward
return self.module(batch_list[0])
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/common/utils.py", line 135, in cls_method
return f(self, *args, **kwargs)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/models/gemnet_oc/gemnet_oc.py", line 1258, in forward
) = self.get_bases(
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/models/gemnet_oc/gemnet_oc.py", line 1129, in get_bases
basis_rad_sph_qint_raw, basis_sph_qint_raw = self.sbf_basis_qint(
File "/home/josegaut/.conda/envs/ocp-models/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/models/gemnet_oc/layers/spherical_basis.py", line 136, in forward
sph_basis = self.spherical_basis(cosφ_cab, θ_cabd)
File "/home/josegaut/PythonModules/ocp-main/ocpmodels/models/gemnet_oc/layers/spherical_basis.py", line 117, in
self.spherical_basis = lambda cosφ, ϑ: (
RuntimeError: cannot reshape tensor of 0 elements into shape [0, -1] because the unspecified dimension size -1 can be any value and is ambiguous

---

The calculation then breaks. It seems to run as expected when using the GemNet-dT checkpoint but thought I would raise this issue in case it's not intended behavior. Here is the full python script I am running:

import numpy as np
from ase.build import add_adsorbate, fcc100, molecule
from ase.constraints import FixAtoms
from ase.optimize import BFGS
from ocpmodels.common.relaxation.ase_utils import OCPCalculator

# Construct a sample structure, similar to the EMT relaxation example!
adslab = read('init.traj')

# Define the calculator
# checkpoint_path = '/home/josegaut/PythonModules/ocp-main/checkpoints/gndt_oc22_all_s2ef.pt'
checkpoint_path = '/home/josegaut/PythonModules/ocp-main/checkpoints/gnoc_finetune_all_s2ef.pt'
calc = OCPCalculator(checkpoint=checkpoint_path)

# Set up the calculator
adslab.set_calculator(calc)

# os.makedirs("data", exist_ok=True)
opt = BFGS(adslab, trajectory='qn.traj',logfile='qn.log')

opt.run(fmax=0.05, steps=100)```

Thanks!

[abhshkdz]

Hey @gauthierlab, all of the other models are GemNet-OC variants, so it's likely a tagging issue -- GemNet-OC uses only the atoms tagged as 1 and 2 to compute quadruplets (see code). In OC20 parlance, tag==1 refers to surface atoms and tag==2 refers to adsorbate atoms. If all the atoms are tagged as 0 (check atoms.get_tags()), no quadruplets are computed, and part of the GemNet-OC forward pass fails. Having some atoms tagged as 1 or 2 in your structure should fix it. (We should probably update the code to throw a nicer error message when the tags are incorrect.)

[gauthierlab]

This was it -- thanks!