code: def NT_Xent(z1, z2, tau=0.5, norm=True):
      """
      Args:
          z1, z2: Tensor of shape [batch_size, z_dim]
          tau: Float. Usually in (0,1].
          norm: Boolean. Whether to apply normlization.
      """
  
      batch_size, _ = z1.size()
      sim_matrix = torch.einsum("ik,jk->ij", z1, z2)
  
      if norm:
          z1_abs = z1.norm(dim=1)
          z2_abs = z2.norm(dim=1)
          sim_matrix = sim_matrix / torch.einsum(
              "i,j->ij", z1_abs, z2_abs
          )
  
      sim_matrix = torch.exp(sim_matrix / tau)
      pos_sim = sim_matrix[range(batch_size), range(batch_size)]
      loss = pos_sim / (sim_matrix.sum(dim=1) - pos_sim)
      loss = -torch.log(loss).mean()
      return loss
  
imports: import itertools
  import torch
  import torch.nn.functional
label_desc: Use image node_id 4 for calling a global function with example usage: NT_Xent(z1, z2, tau, norm) and returns: loss