def __init__(
        self,
        num_features: int,
        momentum: float = 0.01,
        track_running_stats: bool = True,
        eps: float = 0,
    ):
        """Initializes the Whitening2d module with the specified parameters.

        Args:
            num_features:
                Number of features in the input.
            momentum:
                Momentum for the running mean and variance.
            track_running_stats:
                If True, tracks the running mean and variance.
            eps:
                Epsilon for numerical stability.

        Raises:
            RuntimeError: If torch.linalg.solve_triangular is not available in the PyTorch installation.
        """

        super(Whitening2d, self).__init__()

        if not _SOLVE_TRIANGULAR_AVAILABLE:
            raise RuntimeError(
                "Whitening2d depends on torch.linalg.solve_triangular which is not "
                "available in your PyTorch installation. Please update to PyTorch 1.11 "
                "or newer."
            )

        self.running_mean: torch.Tensor
        self.running_variance: torch.Tensor
        self.num_features = num_features
        self.momentum = momentum
        self.track_running_stats = track_running_stats
        self.eps = eps

        if self.track_running_stats:
            self.register_buffer(
                "running_mean", torch.zeros([1, self.num_features, 1, 1])
            )
            self.register_buffer("running_variance", torch.eye(self.num_features))
Initializes the Whitening2d module with the specified parameters.

        Args:
            num_features:
                Number of features in the input.
            momentum:
                Momentum for the running mean and variance.
            track_running_stats:
                If True, tracks the running mean and variance.
            eps:
                Epsilon for numerical stability.

        Raises:
            RuntimeError: If torch.linalg.solve_triangular is not available in the PyTorch installation.
        
lightly/loss/wmse_loss.py