###############
PyTorch Helpers
###############

.. autoclass:: graph_pes.utils.nn.PerElementParameter
   :show-inheritance:
   :members:

.. autoclass:: graph_pes.utils.nn.PerElementEmbedding

.. autoclass:: graph_pes.utils.nn.MLPConfig
   :show-inheritance:
   :members:

.. autoclass:: graph_pes.utils.nn.MLP
    :members: