def __init__(
        self,
        embedding_dim: int = 128,
        momentum: float = 0.01,
        eps: float = 0.0,
        track_running_stats: bool = True,
        w_iter: int = 1,
        w_size: int = 256,
        loss_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = norm_mse_loss,
        num_samples: int = 2,
    ):
        """Initializes the WMSELoss module with the specified parameters.

        Parameters as described in [0].

        Args:
            embedding_dim:
                Dimensionality of the embedding.
            momentum:
                Momentum for the running statistics.
            eps:
                Epsilon for the running statistics.
            track_running_stats:
                Whether to track running statistics.
            w_iter:
                Number of iterations for the whitening.
            w_size:
                Sub-batch size to use for whitening.
            loss_fn:
                Loss function to use for the whitening.
            num_samples:
                Number of samples generated by the transforms for each image.

        Raises:
            ValueError: If w_size is less than twice the size of embedding_dim.
        """
        super().__init__()
        self.whitening = Whitening2d(
            num_features=embedding_dim,
            momentum=momentum,
            eps=eps,
            track_running_stats=track_running_stats,
        )
        if embedding_dim * 2 > w_size:
            raise ValueError(
                "w_size should be at least twice the size of embedding_dim to avoid instabiliy"
            )
        self.w_iter = w_iter
        self.w_size = w_size
        self.loss_f = loss_fn
        self.num_samples = num_samples
        self.num_pairs = num_samples * (num_samples - 1) // 2
Initializes the WMSELoss module with the specified parameters.

        Parameters as described in [0].

        Args:
            embedding_dim:
                Dimensionality of the embedding.
            momentum:
                Momentum for the running statistics.
            eps:
                Epsilon for the running statistics.
            track_running_stats:
                Whether to track running statistics.
            w_iter:
                Number of iterations for the whitening.
            w_size:
                Sub-batch size to use for whitening.
            loss_fn:
                Loss function to use for the whitening.
            num_samples:
                Number of samples generated by the transforms for each image.

        Raises:
            ValueError: If w_size is less than twice the size of embedding_dim.
        
lightly/loss/wmse_loss.py