code: def _normalize(
      z_a: torch.Tensor, z_b: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """Helper function to normalize tensors along the batch dimension."""
      combined = torch.stack([z_a, z_b], dim=0)  # Shape: 2 x N x D
      normalized = F.batch_norm(
          combined.flatten(0, 1),
          running_mean=None,
          running_var=None,
          weight=None,
          bias=None,
          training=True,
      ).view_as(combined)
      return normalized[0], normalized[1]
  
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 1 for calling a global function with example usage: _normalize(z_a, z_b) and returns: normalized,  normalized