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 __init__(
      self,
      num_features: int,
      momentum: float = 0.01,
      track_running_stats: bool = True,
      eps: float = 0,
  ):
      super(Whitening2d, self).__init__()
      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)
          )
  
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 1 to create a new Whitening2d object from inherited base classes: nn with example: obj = Whitening2d(num_features, momentum, track_running_stats, eps)