def forward(self, x: Tensor) -> Tensor:
        """Forward pass through KoLeo Loss.

        Args:
            x: Tensor with shape (batch_size, embedding_size).

        Returns:
            Loss value.
        """
        # Normalize the input tensor
        x = functional.normalize(x, p=2, dim=-1, eps=self.eps)

        # Calculate cosine similarity.
        cos_sim = torch.mm(x, x.t())
        cos_sim.fill_diagonal_(-2)

        # Get nearest neighbors.
        nn_idx = cos_sim.argmax(dim=1)
        nn_dist: Tensor = self.pairwise_distance(x, x[nn_idx])

        # Compute the loss
        loss = -(nn_dist + self.eps).log().mean()

        return loss
Forward pass through KoLeo Loss.

        Args:
            x: Tensor with shape (batch_size, embedding_size).

        Returns:
            Loss value.
        
lightly/loss/koleo_loss.py