gnn_tracking.training.ec
========================

.. py:module:: gnn_tracking.training.ec

.. autoapi-nested-parse::

   Lightning module for edge classifier training.



Classes
-------

.. autoapisummary::

   gnn_tracking.training.ec.ECModule


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

.. py:class:: ECModule(*, loss_fct: torch.nn.Module, **kwargs)

   Bases: :py:obj:`gnn_tracking.training.base.TrackingModule`


   Lightning module for edge classifier training.


   .. py:attribute:: loss_fct


   .. py:method:: get_losses(out: dict[str, Any], data: torch_geometric.data.Data) -> torch.Tensor


   .. py:method:: training_step(batch: torch_geometric.data.Data, batch_idx: int) -> torch.Tensor | None


   .. py:method:: validation_step(batch: torch_geometric.data.Data, batch_idx: int)


   .. py:method:: highlight_metric(metric: str) -> bool