deepmd_gnn.autograd

Autograd helpers shared by GNN model wrappers.

Functions

derive_atomic_virial_from_displacement(→ torch.Tensor)

Derive per-atom virials from atom-energy gradients w.r.t. cell strain.

Module Contents

deepmd_gnn.autograd.derive_atomic_virial_from_displacement(atom_energy: torch.Tensor, displacement: torch.Tensor, nloc: int, create_graph: bool) → torch.Tensor

Derive per-atom virials from atom-energy gradients w.r.t. cell strain.