def forward(self, z0: Tensor, z1: Tensor) -> Tensor:
        """Compute the Model-Aware Contrastive Loss (MACL) for a batch of embeddings.

        Args:
            z0:
                First view embeddings
                Shape (batch_size, embedding_size)
            z1:
                Second view embeddings
                Shape (batch_size, embedding_size)

        Returns:
            loss:
                The computed loss value
        """
        # Normalize embeddings
        z0 = F.normalize(z0, dim=-1, p=2)
        z1 = F.normalize(z1, dim=-1, p=2)

        # Concatenate embeddings
        out = torch.cat([z0, z1], dim=0)
        batch_size = z0.shape[0]
        n_samples = len(out)

        # Compute similarity matrix
        cov = out @ out.T

        # Get positive and negative pairs
        mask = cov.new_ones(cov.shape, dtype=torch.bool)
        mask.diagonal()[:] = False
        mask.diagonal(batch_size)[:] = False
        mask.diagonal(-batch_size)[:] = False
        neg = cov.masked_select(mask).view(n_samples, -1)

        # Get positive pairs from upper and lower diagonals
        u_b = torch.diag(cov, batch_size)
        l_b = torch.diag(cov, -batch_size)
        pos = torch.cat([u_b, l_b], dim=0).reshape(n_samples, 1)

        # Calculate model-aware temperature
        A = torch.mean(pos.detach())
        t = self.temperature * (1 + self.alpha * (A - self.A_0))

        # 6) Compute stable log_softmax
        logits = torch.cat([pos, neg], dim=1) / t
        log_prob = F.log_softmax(logits, dim=1)

        # 7) Extract log probability of the positive pair
        log_pos_prob = log_prob[:, 0]

        # 8) Compute P and V
        P = log_pos_prob.exp()
        V = 1.0 / (1.0 - P + self.eps)  # add eps to avoid division-by-zero

        # 9) Compute final loss
        loss: Tensor = -V.detach() * log_pos_prob

        return loss.mean()
Compute the Model-Aware Contrastive Loss (MACL) for a batch of embeddings.

        Args:
            z0:
                First view embeddings
                Shape (batch_size, embedding_size)
            z1:
                Second view embeddings
                Shape (batch_size, embedding_size)

        Returns:
            loss:
                The computed loss value
        
lightly/loss/macl_loss.py