deepmd_gnn.autograd
===================

.. py:module:: deepmd_gnn.autograd

.. autoapi-nested-parse::

   Autograd helpers shared by GNN model wrappers.



Functions
---------

.. autoapisummary::

   deepmd_gnn.autograd.derive_atomic_virial_from_displacement


Module Contents
---------------

.. py:function:: 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.