gnn_tracking.utils.torch_utils

Utility functions for general torch stuff.

Functions

freeze(→ torch.nn.Module)	Freezes all parameters of a model.
freeze_if(→ torch.nn.Module | None)	Freezes all parameters of a model if do_freeze is True. If model is None,

Module Contents

gnn_tracking.utils.torch_utils.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).

gnn_tracking.utils.torch_utils.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).