def _get_pseudo_labels(
        self, out0: Tensor, out1: Tensor, negatives: Union[Tensor, None] = None
    ) -> Tensor:
        """Computes the soft pseudo labels across negative samples.

        Args:
            out0:
                Output projections of the first set of transformed images (query).
                Shape: bsz x n_ftrs
            out1:
                Output projections of the second set of transformed images (positive sample).
                Shape: bsz x n_ftrs
            negatives:
                Negative samples to compare against. If this is None, the second
                batch of images will be used as negative samples.
                Shape: memory_bank_size x n_ftrs

        Returns:
            Log probability that a positive samples will classify each negative
            sample as the positive sample.
            Shape: bsz x (bsz - 1) or bsz x memory_bank_size
        """
        batch_size, _ = out0.shape
        if negatives is None:
            # Use second batch as negative samples
            # l_pos has shape bsz x 1 and l_neg has shape bsz x bsz
            l_pos = torch.einsum("nc,nc->n", [out0, out1]).unsqueeze(-1)
            l_neg = torch.einsum("nc,ck->nk", [out0, out1.t()])

            # Remove elements on the diagonal
            # l_neg has shape bsz x (bsz - 1)
            l_neg = l_neg.masked_select(
                ~torch.eye(batch_size, dtype=torch.bool, device=l_neg.device)
            ).view(batch_size, batch_size - 1)
        else:
            # Use memory bank as negative samples
            # l_pos has shape bsz x 1 and l_neg has shape bsz x memory_bank_size
            negatives = negatives.to(out0.device)
            l_pos = torch.einsum("nc,nc->n", [out0, out1]).unsqueeze(-1)
            l_neg = torch.einsum("nc,ck->nk", [out0, negatives.clone().detach()])

        # Concatenate such that positive samples are at index 0
        logits = torch.cat([l_pos, l_neg], dim=1)
        # Divide by temperature
        logits = logits / self.t_consistency

        # The input to kl_div is expected to be log(p)
        return torch.nn.functional.log_softmax(logits, dim=-1)
Computes the soft pseudo labels across negative samples.

        Args:
            out0:
                Output projections of the first set of transformed images (query).
                Shape: bsz x n_ftrs
            out1:
                Output projections of the second set of transformed images (positive sample).
                Shape: bsz x n_ftrs
            negatives:
                Negative samples to compare against. If this is None, the second
                batch of images will be used as negative samples.
                Shape: memory_bank_size x n_ftrs

        Returns:
            Log probability that a positive samples will classify each negative
            sample as the positive sample.
            Shape: bsz x (bsz - 1) or bsz x memory_bank_size
        
lightly/loss/regularizer/co2.py