def __init__(
        self,
        beta: float = 0.9,
        rho: float = 8.0,
        gather_distributed: bool = False,
    ):
        """Initializes the TiCoLoss module with the specified parameters.

        Args:
            beta:
                Coefficient for the EMA update of the covariance.
            rho:
                Weight for the covariance term of the loss.
            gather_distributed:
                If True, the cross-correlation matrices from all GPUs are gathered
                    and summed before the loss calculation. Default is False.

        Raises:
            ValueError: If gather_distributed is True but torch.distributed is not available.
        """
        super(TiCoLoss, self).__init__()
        if gather_distributed and not dist.is_available():
            raise ValueError(
                "gather_distributed is True but torch.distributed is not available. "
                "Please set gather_distributed=False or install a torch version with "
                "distributed support."
            )

        self.beta = beta
        self.rho = rho
        self.C: Union[Tensor, None] = None
        self.gather_distributed = gather_distributed
Initializes the TiCoLoss module with the specified parameters.

        Args:
            beta:
                Coefficient for the EMA update of the covariance.
            rho:
                Weight for the covariance term of the loss.
            gather_distributed:
                If True, the cross-correlation matrices from all GPUs are gathered
                    and summed before the loss calculation. Default is False.

        Raises:
            ValueError: If gather_distributed is True but torch.distributed is not available.
        
lightly/loss/tico_loss.py