:py:mod:`gnn_tracking.utils.torch_utils`
========================================

.. py:module:: gnn_tracking.utils.torch_utils

.. autoapi-nested-parse::

   Utility functions for general torch stuff.



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


Functions
~~~~~~~~~

.. autoapisummary::

   gnn_tracking.utils.torch_utils.freeze
   gnn_tracking.utils.torch_utils.freeze_if



.. py:function:: freeze(model: torch.nn.Module) -> torch.nn.Module

   Freezes all parameters of a model.

   :returns: The model with all parameters frozen (but model is also modified in-place).


.. py:function:: freeze_if(model: torch.nn.Module | None, do_freeze: bool = False) -> torch.nn.Module | None

   Freezes all parameters of a model if `do_freeze` is True. If model is None,
   None is returned. This is a trivial convenience function to avoid if-else
   statements.

   :returns: The model with all parameters frozen (but model is also modified in-place).