class_docstr: A custom loss module for CLIP (Contrastive Language–Image Pretraining) training.
  
  This module is specifically designed for calculating the loss in CLIP model training,
  supporting features like local loss calculation, gradient gathering, and label caching
  for efficiency in a distributed training setup.
  
  Parameters:
  local_loss (bool, optional): If True, calculates loss locally on each data parallel process.
                               Defaults to False.
  gather_with_grad (bool, optional): If True, gathers gradients during loss calculation.
                                     Currently not functional in the latest PyTorch version.
                                     Defaults to False.
  cache_labels (bool, optional): If True, caches labels for reuse in subsequent iterations,
                                 improving performance. Defaults to False.
  
  Attributes:
  world_size (int): The size of the data parallel group (number of processes).
  rank (int): The rank of the current process within the data parallel group.
  
  Methods:
  forward(output_tensor): Computes the loss given the model's output tensor. This involves
                          gathering features across processes, computing logits, and
                          calculating the final cross-entropy loss.
class_docstr_tok: ["A","custom","loss","module","for","CLIP","(","Contrastive","Language\u2013Image","Pretraining",")","training",".","This","module","is","specifically","designed","for","calculating","the","loss","in","CLIP","model","training",",","supporting","features","like","local","loss","calculation",",","gradient","gathering",",","and","label","caching","for","efficiency","in","a","distributed","training","setup",".","Parameters",":","local_loss","(","bool",",","optional",")",":","If","True",",","calculates","loss","locally","on","each","data","parallel","process",".","Defaults","to","False",".","gather_with_grad","(","bool",",","optional",")",":","If","True",",","gathers","gradients","during","loss","calculation",".","Currently","not","functional","in","the","latest","PyTorch","version",".","Defaults","to","False",".","cache_labels","(","bool",",","optional",")",":","If","True",",","caches","labels","for","reuse","in","subsequent","iterations",",","improving","performance",".","Defaults","to","False",".","Attributes",":","world_size","(","int",")",":","The","size","of","the","data","parallel","group","(","number","of","processes",")",".","rank","(","int",")",":","The","rank","of","the","current","process","within","the","data","parallel","group",".","Methods",":","forward","(","output_tensor",")",":","Computes","the","loss","given","the","model","'s","output","tensor",".","This","involves","gathering","features","across","processes",",","computing","logits",",","and","calculating","the","final","cross-entropy","loss","."]
code: def forward(self, output_tensor):
      image_features, text_features, logit_scale = output_tensor
      device = image_features.device
      if self.world_size > 1:
          all_image_features, all_text_features = gather_features(
              image_features,
              text_features,
              self.local_loss,
              self.gather_with_grad,
          )
  
          if self.local_loss:
              logits_per_image = (
                  logit_scale * image_features @ all_text_features.T
              )
              logits_per_text = (
                  logit_scale * text_features @ all_image_features.T
              )
          else:
              logits_per_image = (
                  logit_scale * all_image_features @ all_text_features.T
              )
              logits_per_text = logits_per_image.T
      else:
          logits_per_image = (
              logit_scale * image_features @ text_features.T
          )
          logits_per_text = (
              logit_scale * text_features @ image_features.T
          )
  
      # calculated ground-truth and cache if enabled
      num_logits = logits_per_image.shape[0]
      if (
          self.prev_num_logits != num_logits
          or device not in self.labels
      ):
          labels = torch.arange(
              num_logits, device=device, dtype=torch.long
          )
          if self.world_size > 1 and self.local_loss:
              labels = labels + num_logits * self.rank
          if self.cache_labels:
              self.labels[device] = labels
              self.prev_num_logits = num_logits
      else:
          labels = self.labels[device]
  
      total_loss = (
          F.cross_entropy(logits_per_image, labels)
          + F.cross_entropy(logits_per_text, labels)
      ) / 2
  
      # TODO (yuya): this is not necessary; not necessary if global!
      reduced_loss = average_losses_across_data_parallel_group(
          [total_loss]
      )
      return total_loss, {"loss": reduced_loss}
  
file_path: NeMo/nemo/collections/multimodal/losses/clip_loss.py
imports: import torch
  import torch.distributed.nn
  import torch.nn
  from torch import distributed
  from torch.nn import functional
  from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group
label_desc: Use image node_id 2 for calling the ClipLoss obj's underlying member method code with example usage: obj.forward(output_tensor) and returns: total_loss,  dict