How to use OC20 IS2RE Dimnet++ checkpoint

[liudean998]

What would you like to report?

Hello,

I am trying to use the OC20 IS2RE DimeNet++ checkpoint you provided to predict the relaxed adsorption energy of the initial structure. However, it seems that the parameters in the checkpoint do not match the current DimeNet++ parameters.

Here is my code:

import torch
from fairchem.core.models.dimenet_plus_plus import DimeNetPlusPlus

gpu_id = 1
device = torch.device(f"cuda:{gpu_id}" if torch.cuda.is_available() else "cpu")

checkpoint = torch.load('/home/Downloads/dimenetpp_all.pt', map_location=device)
config = checkpoint['config']['model_attributes']

model = DimeNetPlusPlus(
hidden_channels=config["hidden_channels"],
out_channels=config["out_channels"],
num_blocks=config["num_blocks"],
int_emb_size=config["int_emb_size"],
basis_emb_size=config["basis_emb_size"],
out_emb_channels=config["out_emb_channels"],
num_spherical=config["num_spherical"],
num_radial=config["num_radial"],
)

model.load_state_dict(checkpoint['state_dict'])
model.to(device)

model.eval()

However, I encountered the following error:
TypeError: DimeNetPlusPlus.__init__() missing 3 required positional arguments: 'out_channels', 'int_emb_size', and 'basis_emb_size'
How can I resolve this issue?
Alternatively, is there a tutorial for quickly performing inference using the IS2RE checkpoint?

Thank you!