code: def gather_features(
      image_features: Tensor,
      text_features: Tensor,
      use_distributed: bool,
  ) -> Tuple[Tensor, Tensor]:
      """
      Helper function that allows us to gather image and text features from all DDP ranks in a differentiable manner
  
      Args:
          image_features: Image features
          text_features: Text features
          use_distributed: DDP training or not
  
      Shapes:
          image_features: Shape is [N, d]
          text_features: Shape is [N, d]
          where N is the local batch size and d is the feature dimension.
  
      Returns:
          A tuple of gathered image and text features across all GPUs. In case of a DDP task, each feature
          tensor has a dimension of [G, d] where G=NW is the effective batch size and W is world
          size (or total number of GPUs).
      """
      if use_distributed:
          # gather features from all ranks
          # [N, d] x W --> [G, d] where N and G=NW are the local and global batch sizes and W is the world size
          gathered_image_features = gather_all_features(
              features=image_features, dim=0
          )
          # [N d] x W --> [G, d]
          gathered_text_features = gather_all_features(
              features=text_features, dim=0
          )
          return gathered_image_features, gathered_text_features
      return image_features, text_features
  
file_path: ml-cvnets/loss_fn/multi_modal_img_text/contrastive_loss_clip.py
imports: import argparse
  from typing import Any, Dict, Tuple
  import torch
  from torch import Tensor
  from torch.nn import functional
  from loss_fn import LOSS_REGISTRY
  from loss_fn.multi_modal_img_text.base_multi_modal_img_text_criteria import BaseMultiModalImageTextCriteria
  from utils import logger
  from utils.tensor_utils import gather_all_features
label_desc: Use image node_id 1 for calling a global function with example usage: gather_features(image_features, text_features, use_distributed) and returns: image_features,  text_features, gathered_image_features,  gathered_text_features