Fix double precision in Distance

torchmd / torchmd-net

Training neural network potentials

MIT License

335 stars 75 forks source link

Fix double precision in Distance #170

Closed RaulPPelaez closed 1 year ago

RaulPPelaez commented 1 year ago

Current Distance module fails when loop=True and positions have dtype=torch.double. This PR fixes this by making sure edge_weight has the same dtype as edge_vectors in this specific case.