Neural networks that feel atomic forces, now with a v2 rewrite
A PyTorch library for training ANI-style neural network potentials, recently overhauled with breaking changes and C++/CUDA extensions.

What it does
TorchANI trains and runs “ANI-style” neural network interatomic potentials — essentially machine-learned approximations of quantum mechanical forces between atoms. You feed it molecular geometries; it predicts energies and forces fast enough to drive molecular dynamics simulations. The project is maintained by the Roitberg group and has been around long enough to need a migration guide for its own 2.0 rewrite.
The interesting bit
The library ships with custom C++ and CUDA extensions for descriptor computation and network inference, which you build post-install with an ani build-extensions command. It also exposes a command-line interface (ani --help) and plugs into Amber for mixed quantum/classical (ML/MM) simulations via a separate interface project.
Key highlights
- Requires PyTorch ≥ 2.0; tested against PyTorch 2.8 and CUDA 12.8
- Custom CUDA/C++ extensions for GPU acceleration; CPU-only runs are explicitly “degraded”
- CLI utility
aniincluded for common operations - Migration guide provided for 1.x users; legacy state dicts accessible via
.legacy_state_dict() - Conda package exists but is unmaintained; pip is the recommended install path even inside conda envs
Caveats
- Untested on AMD GPUs (ROCm/HIP) and on Apple Metal Performance Shaders
- macOS installation requires manual tweaks to
environment.yaml; no CUDA support - Conda packaging is effectively deprecated and CI-gated behind specific branch name strings
Verdict
Computational chemists and molecular dynamics hackers already using PyTorch should look here; if you’re not simulating atoms, this is specialized tooling with a nontrivial install path.