def _loss(
        self, out0: Tensor, out1: Tensor, out0_all: Tensor, out1_all: Tensor
    ) -> Tensor:
        """Calculates DCL loss for out0 with respect to its positives in out1
        and the negatives in out1, out0_all, and out1_all.

        This code implements Equation 6 in [0], including the sum over all images `i`
        but with `k` fixed at 0.

        Args:
            out0:
                Output projections of the first set of transformed images.
                Shape: (batch_size, embedding_size)
            out1:
                Output projections of the second set of transformed images.
                Shape: (batch_size, embedding_size)
            out0_all:
                Output projections of the first set of transformed images from
                all distributed processes/gpus. Should be equal to out0 in an
                undistributed setting.
                Shape: (batch_size * world_size, embedding_size)
            out1_all:
                Output projections of the second set of transformed images from
                all distributed processes/GPUs. Should be equal to out1 in an
                undistributed setting.
                Shape: (batch_size * world_size, embedding_size)

        Returns:
            Mean loss over the mini-batch.
        """
        # Create diagonal mask that only selects similarities between
        # representations of the same images
        batch_size = out0.shape[0]
        if self.gather_distributed and dist.world_size() > 1:
            diag_mask = dist.eye_rank(batch_size, device=out0.device)
        else:
            diag_mask = torch.eye(batch_size, device=out0.device, dtype=torch.bool)

        # Calculate similarities (n = batch_size, m = batch_size * world_size)
        sim_00 = torch.einsum("nc,mc->nm", out0, out0_all) / self.temperature
        sim_01 = torch.einsum("nc,mc->nm", out0, out1_all) / self.temperature

        positive_loss = -sim_01[diag_mask]
        if self.weight_fn:
            positive_loss = positive_loss * self.weight_fn(out0, out1)

        # Remove simliarities between same views of the same image
        sim_00 = sim_00[~diag_mask].view(batch_size, -1)

        # Remove similarities between different views of the same images
        # This is the key difference compared to NTXentLoss
        sim_01 = sim_01[~diag_mask].view(batch_size, -1)

        all_negs = torch.cat([sim_00, sim_01], dim=1)
        negative_loss = torch.logsumexp(all_negs, dim=1)  # log(sum exp over *all* negs)
        return (positive_loss + negative_loss).mean()
Calculates DCL loss for out0 with respect to its positives in out1
        and the negatives in out1, out0_all, and out1_all.

        This code implements Equation 6 in [0], including the sum over all images `i`
        but with `k` fixed at 0.

        Args:
            out0:
                Output projections of the first set of transformed images.
                Shape: (batch_size, embedding_size)
            out1:
                Output projections of the second set of transformed images.
                Shape: (batch_size, embedding_size)
            out0_all:
                Output projections of the first set of transformed images from
                all distributed processes/gpus. Should be equal to out0 in an
                undistributed setting.
                Shape: (batch_size * world_size, embedding_size)
            out1_all:
                Output projections of the second set of transformed images from
                all distributed processes/GPUs. Should be equal to out1 in an
                undistributed setting.
                Shape: (batch_size * world_size, embedding_size)

        Returns:
            Mean loss over the mini-batch.
        
lightly/loss/dcl_loss.py