class_docstr: Implementation of the whitening layer as described in [0].
  
  [0] W-MSE, 2021, https://arxiv.org/pdf/2007.06346.pdf
class_docstr_tok: ["Implementation","of","the","whitening","layer","as","described","in","[","0","]",".","[","0","]","W-MSE",",","2021",",","https",":","\/\/arxiv.org\/pdf\/2007.06346.pdf"]
code: def forward(self, x: torch.Tensor) -> torch.Tensor:
      x = x.unsqueeze(2).unsqueeze(3)
      m = (
          x.mean(0)
          .view(self.num_features, -1)
          .mean(-1)
          .view(1, -1, 1, 1)
      )
      if (
          not self.training and self.track_running_stats
      ):  # for inference
          m = self.running_mean
      xn = x - m
  
      T = (
          xn.permute(1, 0, 2, 3)
          .contiguous()
          .view(self.num_features, -1)
      )
      f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)
  
      eye = torch.eye(self.num_features).type(f_cov.type())
  
      if (
          not self.training and self.track_running_stats
      ):  # for inference
          f_cov = self.running_variance
  
      f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye
  
      inv_sqrt = torch.linalg.solve_triangular(
          torch.linalg.cholesky(f_cov_shrinked), eye, upper=False
      )
  
      inv_sqrt = inv_sqrt.contiguous().view(
          self.num_features, self.num_features, 1, 1
      )
  
      decorrelated = F.conv2d(xn, inv_sqrt)
  
      if self.training and self.track_running_stats:
          self.running_mean = torch.add(
              self.momentum * m.detach(),
              (1 - self.momentum) * self.running_mean,
              out=self.running_mean,
          )
          self.running_variance = torch.add(
              self.momentum * f_cov.detach(),
              (1 - self.momentum) * self.running_variance,
              out=self.running_variance,
          )
  
      return decorrelated.squeeze(2).squeeze(2)
  
file_path: lightly/lightly/loss/wmse_loss.py
imports: from typing import Callable
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the Whitening2d obj's underlying member method code with example usage: obj.forward(x) and returns: decorrelated