Raimondas Galvelis

Could you post a table with the raw numbers of your benchmarks?

Yes, the speed up of DHFR and FC9 looks very good.

Ping: @PhilippThoelke @giadefa

Also checking the loss function code and comments, it seems `dy` is interpreted as gradients (not forces). https://github.com/torchmd/torchmd-net/blob/db72e12461b1bce9bb3ddebd093fa11a803d37ca/torchmdnet/module.py#L71-L132

The proof-of-concept kernel is here (https://github.com/torchmd/torchmd-net/pull/61). It was optimized for small molecules, but at the moment neither batching, nor the periodic boundary condition are not supported.

Thanks @peastman! I'll try to move the code to NNPOps and make something working tomorrow.

@peastman you are working just on the optimization of the equivariant transformer, aren't? We also need the graph network, which is effectively the same architecture as SchNet. @PhilippThoelke correct? So,...

This PR is discontinued. The code is being move to NNPOps (https://github.com/openmm/NNPOps/pull/58)

Regarding the interface, it should look and work like this: ```python # Create or load a model in any way model = TorchMD_GN() # Optional: train or do what ever...

For a moment, it seems all the PyTorch-Geometric packages are broken (https://github.com/pyg-team/pytorch_geometric/issues/3660).