code: def init_distributed(cuda):
      """
      Initializes distributed backend.
      :param cuda: (bool) if True initializes nccl backend, if False initializes
          gloo backend
      """
      world_size = int(os.environ.get("WORLD_SIZE", 1))
      distributed = world_size > 1
      if distributed:
          backend = "nccl" if cuda else "gloo"
          torch.distributed.init_process_group(
              backend=backend, init_method="env://"
          )
          assert torch.distributed.is_initialized()
      return distributed
  
file_path: flash-attention/training/src/utils/distributed.py
imports: import os
  from contextlib import contextmanager
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: init_distributed(cuda) and returns: distributed