code: def forward(
      self, z_a: torch.Tensor, z_b: torch.Tensor
  ) -> torch.Tensor:
      # code from https://github.com/facebookresearch/barlowtwins/blob/main/main.py
  
      N = z_a.size(0)
  
      # empirical cross-correlation matrix
      c = self.bn(z_a).T @ self.bn(z_b)
  
      # sum the cross-correlation matrix between all gpus
      c.div_(N)
      if self.gather_distributed:
          torch.distributed.all_reduce(c)
  
      on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
      off_diag = off_diagonal(c).pow_(2).sum()
      loss = on_diag + self.lambda_param * off_diag
      return loss
  
file_path: lightly/tests/loss/test_barlow_twins_loss.py
imports: import pytest
  import torch
  import torch.nn
  from pytest_mock import MockerFixture
  from torch import distributed
  from lightly.loss.barlow_twins_loss import BarlowTwinsLoss
label_desc: Use image node_id 2 for calling the BarlowTwinsLossReference obj's underlying member method code with example usage: obj.forward(z_a, z_b) and returns: loss