def regularization_loss(self, mean_anchor_probs: Tensor) -> Tensor:
        """Calculates regularization loss with a custom target distribution.

        Args:
            mean_anchor_probs:
                The mean anchor probabilities.

        Returns:
            The calculated regularization loss.
        """
        target_dist = self.target_distribution(mean_anchor_probs).to(
            mean_anchor_probs.device
        )
        loss = F.kl_div(
            input=mean_anchor_probs.log(), target=target_dist, reduction="sum"
        )
        return loss
Calculates regularization loss with a custom target distribution.

        Args:
            mean_anchor_probs:
                The mean anchor probabilities.

        Returns:
            The calculated regularization loss.
        
lightly/loss/pmsn_loss.py