Error when using pytorch lightning

I am trying to define an ANI model along with the AEVComputer (with cuda enabled) module within a Pytorch Lightning Module, but I am getting the following error:

RuntimeError: coordinates, species, and aev_params should be on the same device

I have seen that some of the parameters are registered as buffers, but some are not. Please let me know what you think.

Kev

[yueyericardo]

Hi, the error came from the openmm-ml wrapper. A temp fixed version work ONLY for GPU could be found at: https://github.com/yueyericardo/openmm-ml/commit/1d1d3f24f40becdcd8a36431c8d0900d98eb1304#diff-911692ca194bf903c77d038662969ad3277dcf2fa8b3b3048d95a5aa3af59de1

It is using cuaev use_cuda_extension for aev calculation, but it currently does not support pbc, so if you want to use cuaev, you have to change your script slightly to

pdb = PDBFile("aaa.pdb")
# add this line
pdb.topology.setPeriodicBoxVectors(None)
potential = MLPotential('ani2x')

Our internal version has some other updates to make it faster, but it currently is not open source yet. In the meanwhile, openmm team is building NNPOPS for ani and schnet, you could track the progress here Add example of using NNPOps with openmm-torch?!

Edit: BTW, our conda-forge package includes the latest public build with cuaev: you could install it directly by

conda install -c conda-forge torchani

[kryczko]

I am still getting the same issue I showed above while using an ANI model within pytorch lightning. Any ideas how to fix it?