Use pytorch geometric?

[lilyminium]

DGL is no longer very up-to-date and I believe pyg will give us torchscript.

This will be non-trivial to do currently due to the somewhat hodgepodge way NAGL got updated to allow non-dgl inference. Possibly the new espaloma architecture will be set up in a way for us to move onto that easily?

[mattwthompson]

https://docs.nvidia.com/physicsnemo/latest/resources/dgl_to_pyg_migration.html

[mattwthompson]

Possible brige here? https://pytorch-geometric.readthedocs.io/en/2.5.3/modules/utils.html#torch_geometric.utils.from_dgl

For intermediate testing and validation, however, just using NetworkX graphs may be easier