def forward(
        self,
        z_a: torch.Tensor,
        z_b: torch.Tensor,
        update_covariance_matrix: bool = True,
    ) -> torch.Tensor:
        """Computes the TiCo loss.

        It maximizes the agreement among embeddings of different distorted versions of the same image
        while avoiding collapse using Covariance matrix.

        Args:
            z_a:
                Tensor of shape [batch_size, num_features=256]. Output of the learned backbone.
            z_b:
                Tensor of shape [batch_size, num_features=256]. Output of the momentum updated backbone.
            update_covariance_matrix:
                Parameter to update the covariance matrix at each iteration.

        Returns:
            The computed loss.

        Raises:
            AssertionError: If z_a or z_b have a batch size <= 1.
            AssertionError: If z_a and z_b do not have the same shape.
        """

        assert (
            z_a.shape[0] > 1 and z_b.shape[0] > 1
        ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}"
        assert (
            z_a.shape == z_b.shape
        ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}."

        # gather all batches
        if self.gather_distributed and dist.is_initialized():
            world_size = dist.get_world_size()
            if world_size > 1:
                z_a = torch.cat(gather(z_a), dim=0)
                z_b = torch.cat(gather(z_b), dim=0)

        # Normalize image
        z_a = torch.nn.functional.normalize(z_a, dim=1)
        z_b = torch.nn.functional.normalize(z_b, dim=1)

        # Compute auxiliary matrix B
        B = torch.mm(z_a.T, z_a).detach() / z_a.shape[0]

        # Initialize covariance matrix
        if self.C is None:
            self.C = B.new_zeros(B.shape).detach()

        # Compute loss
        C = self.beta * self.C + (1 - self.beta) * B

        transformative_invariance_loss = 1.0 - (z_a * z_b).sum(dim=1).mean()
        covariance_contrast_loss = self.rho * (torch.mm(z_a, C) * z_a).sum(dim=1).mean()

        loss: Tensor = transformative_invariance_loss + covariance_contrast_loss

        # Update covariance matrix
        if update_covariance_matrix:
            self.C = C.detach()

        return loss
Computes the TiCo loss.

        It maximizes the agreement among embeddings of different distorted versions of the same image
        while avoiding collapse using Covariance matrix.

        Args:
            z_a:
                Tensor of shape [batch_size, num_features=256]. Output of the learned backbone.
            z_b:
                Tensor of shape [batch_size, num_features=256]. Output of the momentum updated backbone.
            update_covariance_matrix:
                Parameter to update the covariance matrix at each iteration.

        Returns:
            The computed loss.

        Raises:
            AssertionError: If z_a or z_b have a batch size <= 1.
            AssertionError: If z_a and z_b do not have the same shape.
        
lightly/loss/tico_loss.py