class_docstr: Implementation of the loss described in 'Whitening for
  Self-Supervised Representation Learning' [0].
  
  - [0] W-MSE, 2021, https://arxiv.org/pdf/2007.06346.pdf
  
  Examples:
      >>> # initialize loss function
      >>> loss_fn = WMSELoss(num_samples=2)
      >>> transform_fn = WMSETransform(num_samples=2)
      >>>
      >>> # generate the transformed samples
      >>> samples = transform_fn(image)
      >>>
      >>> # feed through encoder head
      >>> h = torch.cat([model(s) for s in samples])
      >>>
      >>> # calculate loss
      >>> loss = loss_fn(samples, h)
class_docstr_tok: ["Implementation","of","the","loss","described","in","'Whitening","for","Self-Supervised","Representation","Learning","'","[","0","]",".","-","[","0","]","W-MSE",",","2021",",","https",":","\/\/arxiv.org\/pdf\/2007.06346.pdf","Examples",":",">",">",">","#","initialize","loss","function",">",">",">","loss_fn","=","WMSELoss","(","num_samples=2",")",">",">",">","transform_fn","=","WMSETransform","(","num_samples=2",")",">",">",">",">",">",">","#","generate","the","transformed","samples",">",">",">","samples","=","transform_fn","(","image",")",">",">",">",">",">",">","#","feed","through","encoder","head",">",">",">","h","=","torch.cat","(","[","model","(","s",")","for","s","in","samples","]",")",">",">",">",">",">",">","#","calculate","loss",">",">",">","loss","=","loss_fn","(","samples",",","h",")"]
docstr: 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.
docstr_tok: ["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","."]
code: 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
  
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 WMSELoss obj's underlying member method code with example usage: obj.forward(input) and returns: loss