class_docstr: Implementation of the Barlow Twins Loss from Barlow Twins[0] paper.
  This code specifically implements the Figure Algorithm 1 from [0].
  
  [0] Zbontar,J. et.al, 2021, Barlow Twins... https://arxiv.org/abs/2103.03230
  
      Examples:
  
      >>> # initialize loss function
      >>> loss_fn = BarlowTwinsLoss()
      >>>
      >>> # generate two random transforms of images
      >>> t0 = transforms(images)
      >>> t1 = transforms(images)
      >>>
      >>> # feed through SimSiam model
      >>> out0, out1 = model(t0, t1)
      >>>
      >>> # calculate loss
      >>> loss = loss_fn(out0, out1)
code: def forward(
      self, z_a: torch.Tensor, z_b: torch.Tensor
  ) -> torch.Tensor:
      # normalize repr. along the batch dimension
      z_a_norm, z_b_norm = _normalize(z_a, z_b)
  
      N = z_a.size(0)
  
      # cross-correlation matrix
      c = z_a_norm.T @ z_b_norm
      c.div_(N)
  
      # sum cross-correlation matrix between multiple gpus
      if self.gather_distributed and dist.is_initialized():
          world_size = dist.get_world_size()
          if world_size > 1:
              c = c / world_size
              dist.all_reduce(c)
  
      invariance_loss = torch.diagonal(c).add_(-1).pow_(2).sum()
      redundancy_reduction_loss = _off_diagonal(c).pow_(2).sum()
      loss = (
          invariance_loss
          + self.lambda_param * redundancy_reduction_loss
      )
  
      return loss
  
file_path: lightly/lightly/loss/barlow_twins_loss.py
imports: from typing import Tuple
  import torch
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the BarlowTwinsLoss obj's underlying member method code with example usage: obj.forward(z_a, z_b) and returns: loss