Using TorchMD_ET to make node-level predictions

[arunraja-hub]

Hi, I am trying to use TorchMD_ET to make node-level predictions like partial charges by removing the aggregation in https://github.com/torchmd/torchmd-net/blob/main/torchmdnet/models/torchmd_et.py. It works and the MSE loss does reduce from epoch to epoch but it is not as performant as a simple EGNN. For the final node level predictions I am using a simple MLP with the node embeddings from TorchMD as the input. Is there some changes I could do to improve the node level performance or was TorchMD purpose-built for molecule-level properties?

[dav0dea]

Hi, accuracy and speed of the model depend on the chosen architecture (we support these) and configuration of the model (e.g. embedding dimension, number of layers, ...). You may have to play around with the hyperparameters to find a configuration that yields optimal performance for your task. While the example configurations we provide (here) were optimized for molecule-level predictions, the architectures are quite general and can be used for a range of different tasks.