def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Calculates the W-MSE loss.

        Args:
            input: Tensor with shape (batch_size * num_samples, embedding_dim).

        Returns:
            Aggregate W-MSE loss over all sub-batches.

        Raises:
            RuntimeError: If the batch size is not divisible by num_samples.
            ValueError: If the batch size is smaller than w_size.
        """
        if input.shape[0] % self.num_samples != 0:
            raise RuntimeError("input batch size must be divisible by num_samples")

        bs = input.shape[0] // self.num_samples

        if bs < self.w_size:
            raise ValueError("batch size must be greater than or equal to w_size")
        loss = torch.tensor(0.0, device=input.device, requires_grad=True)

        for _ in range(self.w_iter):
            z = torch.empty_like(input)
            perm = torch.randperm(bs).view(-1, self.w_size)
            for idx in perm:
                for i in range(self.num_samples):
                    z[idx + i * bs] = self.whitening(input[idx + i * bs])
            for i in range(self.num_samples - 1):
                for j in range(i + 1, self.num_samples):
                    x0 = z[i * bs : (i + 1) * bs]
                    x1 = z[j * bs : (j + 1) * bs]
                    loss = loss + self.loss_f(x0, x1)
        loss = loss / (self.w_iter * self.num_pairs)
        return loss
Calculates the W-MSE loss.

        Args:
            input: Tensor with shape (batch_size * num_samples, embedding_dim).

        Returns:
            Aggregate W-MSE loss over all sub-batches.

        Raises:
            RuntimeError: If the batch size is not divisible by num_samples.
            ValueError: If the batch size is smaller than w_size.
        
lightly/loss/wmse_loss.py