DimeNet++ Issue - Returning zeros in output tensor?

[srikarkampalli]

Hello everyone, so I was experimenting with various GNNs related to chemistry, and came across both DimeNet and DimeNet++, which were thankfully implemented in PyTorch Geometric, so I decided to test them out by passing atomic lists and coordinate tensors through both DimeNet and DimeNet++.

Once I processed the lists and passed them through DimeNet++ (in this example), all I got was an output Tensor with zeros for each element in the tensor. Is this on purpose, or am I doing something wrong with the code? I already had PyG (nightly), torch-cluster, torch-scatter, and torch-sparse in my runtime, also, if that information helps.

By the way, I do know that a PR was made for this.

The code implementation is below

import torch
from rdkit import Chem
from rdkit.Chem import AllChem

# Load molecule and add hydrogens
smiles = "CCO"  # ethanol
mol = Chem.MolFromSmiles(smiles)
mol = Chem.AddHs(mol)

# coords
AllChem.EmbedMolecule(mol)
AllChem.UFFOptimizeMolecule(mol)

# get z and pos list
z = torch.tensor([atom.GetAtomicNum() for atom in mol.GetAtoms()], dtype=torch.long)

conf = mol.GetConformer()
coords = [conf.GetAtomPosition(i) for i in range(mol.GetNumAtoms())]
pos = torch.tensor([[pt.x, pt.y, pt.z] for pt in coords], dtype=torch.float)

from torch_geometric.nn.models import DimeNetPlusPlus

model = DimeNetPlusPlus(
  hidden_channels=128,
  out_channels=5,
  num_blocks=4,
  int_emb_size=64,
  basis_emb_size=8,
  out_emb_channels=256,
  num_spherical=7,
  num_radial=6,
  cutoff=5.0,
  max_num_neighbors=32,
  envelope_exponent=5,
  num_before_skip=1,
  num_after_skip=2,
  num_output_layers=3,
)

print(model(z, pos))

And here is the output:

tensor([0., 0., 0., 0., 0.], grad_fn=<SumBackward1>)