class_docstr: Implementation of the VICReg loss [0].
  
  This implementation is based on the code published by the authors [1].
  
  - [0] VICReg, 2022, https://arxiv.org/abs/2105.04906
  - [1] https://github.com/facebookresearch/vicreg/
  
  Attributes:
      lambda_param:
          Scaling coefficient for the invariance term of the loss.
      mu_param:
          Scaling coefficient for the variance term of the loss.
      nu_param:
          Scaling coefficient for the covariance term of the loss.
      gather_distributed:
          If True then the cross-correlation matrices from all gpus are gathered and
          summed before the loss calculation.
      eps:
          Epsilon for numerical stability.
  
  Examples:
  
      >>> # initialize loss function
      >>> loss_fn = VICRegLoss()
      >>>
      >>> # generate two random transforms of images
      >>> t0 = transforms(images)
      >>> t1 = transforms(images)
      >>>
      >>> # feed through model
      >>> out0, out1 = model(t0, t1)
      >>>
      >>> # calculate loss
      >>> loss = loss_fn(out0, out1)
class_docstr_tok: ["Implementation","of","the","VICReg","loss","[","0","]",".","This","implementation","is","based","on","the","code","published","by","the","authors","[","1","]",".","-","[","0","]","VICReg",",","2022",",","https",":","\/\/arxiv.org\/abs\/2105.04906","-","[","1","]","https",":","\/\/github.com\/facebookresearch\/vicreg\/","Attributes",":","lambda_param",":","Scaling","coefficient","for","the","invariance","term","of","the","loss",".","mu_param",":","Scaling","coefficient","for","the","variance","term","of","the","loss",".","nu_param",":","Scaling","coefficient","for","the","covariance","term","of","the","loss",".","gather_distributed",":","If","True","then","the","cross-correlation","matrices","from","all","gpus","are","gathered","and","summed","before","the","loss","calculation",".","eps",":","Epsilon","for","numerical","stability",".","Examples",":",">",">",">","#","initialize","loss","function",">",">",">","loss_fn","=","VICRegLoss","(",")",">",">",">",">",">",">","#","generate","two","random","transforms","of","images",">",">",">","t0","=","transforms","(","images",")",">",">",">","t1","=","transforms","(","images",")",">",">",">",">",">",">","#","feed","through","model",">",">",">","out0",",","out1","=","model","(","t0",",","t1",")",">",">",">",">",">",">","#","calculate","loss",">",">",">","loss","=","loss_fn","(","out0",",","out1",")"]
docstr: Returns VICReg loss.
  
  Args:
      z_a:
          Tensor with shape (batch_size, ..., dim).
      z_b:
          Tensor with shape (batch_size, ..., dim).
docstr_tok: ["Returns","VICReg","loss",".","Args",":","z_a",":","Tensor","with","shape","(","batch_size",",","...",",","dim",")",".","z_b",":","Tensor","with","shape","(","batch_size",",","...",",","dim",")","."]
code: def forward(
      self, z_a: torch.Tensor, z_b: torch.Tensor
  ) -> torch.Tensor:
      """Returns VICReg loss.
  
      Args:
          z_a:
              Tensor with shape (batch_size, ..., dim).
          z_b:
              Tensor with shape (batch_size, ..., dim).
      """
      assert (
          z_a.shape[0] > 1 and z_b.shape[0] > 1
      ), f"z_a and z_b must have batch size > 1 but found {z_a.shape[0]} and {z_b.shape[0]}"
      assert (
          z_a.shape == z_b.shape
      ), f"z_a and z_b must have same shape but found {z_a.shape} and {z_b.shape}."
  
      # invariance term of the loss
      inv_loss = invariance_loss(x=z_a, y=z_b)
  
      # gather all batches
      if self.gather_distributed and dist.is_initialized():
          world_size = dist.get_world_size()
          if world_size > 1:
              z_a = torch.cat(gather(z_a), dim=0)
              z_b = torch.cat(gather(z_b), dim=0)
  
      var_loss = 0.5 * (
          variance_loss(x=z_a, eps=self.eps)
          + variance_loss(x=z_b, eps=self.eps)
      )
      cov_loss = covariance_loss(x=z_a) + covariance_loss(x=z_b)
  
      loss = (
          self.lambda_param * inv_loss
          + self.mu_param * var_loss
          + self.nu_param * cov_loss
      )
      return loss
  
file_path: lightly/lightly/loss/vicreg_loss.py
imports: import torch
  import torch.distributed
  import torch.nn.functional
  from torch import Tensor
  from lightly.utils.dist import gather
label_desc: Use image node_id 2 for calling the VICRegLoss obj's underlying member method code with example usage: obj.forward(z_a, z_b) and returns: loss