def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
        """Computes the Barlow Twins loss for the given projections.

        Args:
            z_a: Output projections of the first set of transformed images.
            z_b: Output projections of the second set of transformed images.

        Returns:
            Computed Barlow Twins Loss.
        """

        # Normalize repr. along the batch dimension
        z_a_norm, z_b_norm = _normalize(z_a), _normalize(z_b)

        N = z_a.size(0)

        # Compute the cross-correlation matrix
        c = z_a_norm.T @ z_b_norm
        c.div_(N)

        # Aggregate and normalize the cross-correlation matrix between multiple GPUs
        if self.gather_distributed and dist.is_initialized():
            world_size = dist.get_world_size()
            if world_size > 1:
                c = c / world_size
                dist.all_reduce(c)

        invariance_loss = torch.diagonal(c).add_(-1).pow_(2).sum()
        redundancy_reduction_loss = _off_diagonal(c).pow_(2).sum()
        loss: Tensor = invariance_loss + self.lambda_param * redundancy_reduction_loss

        return loss
Computes the Barlow Twins loss for the given projections.

        Args:
            z_a: Output projections of the first set of transformed images.
            z_b: Output projections of the second set of transformed images.

        Returns:
            Computed Barlow Twins Loss.
        
lightly/loss/barlow_twins_loss.py