text: import torch
  import torch.distributed as dist
  
  class BarlowTwinsLoss(torch.nn.Module):
      """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)
  
      """
  
      def __init__(
          self, 
          lambda_param: float = 5e-3, 
          gather_distributed : bool = False
      ):
          """Lambda param configuration with default value like in [0]
  
          Args:
              lambda_param: 
                  Parameter for importance of redundancy reduction term. 
                  Defaults to 5e-3 [0].
              gather_distributed:
                  If True then the cross-correlation matrices from all gpus are 
                  gathered and summed before the loss calculation.
          """
          super(BarlowTwinsLoss, self).__init__()
          self.lambda_param = lambda_param
          self.gather_distributed = gather_distributed
  
      def forward(self, z_a: torch.Tensor, z_b: torch.Tensor) -> torch.Tensor:
  
          device = z_a.device
  
          # normalize repr. along the batch dimension
          z_a_norm = (z_a - z_a.mean(0)) / z_a.std(0) # NxD
          z_b_norm = (z_b - z_b.mean(0)) / z_b.std(0) # NxD
  
          N = z_a.size(0)
          D = z_a.size(1)
  
          # cross-correlation matrix
          c = torch.mm(z_a_norm.T, z_b_norm) / N # DxD
  
          # 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)
  
          # loss
          c_diff = (c - torch.eye(D, device=device)).pow(2) # DxD
          # multiply off-diagonal elems of c_diff by lambda
          c_diff[~torch.eye(D, dtype=bool)] *= self.lambda_param
          loss = c_diff.sum()
  
          return loss
  