code: def reduce_loss_dict(loss_dict):
      
      world_size = get_world_size()
      if world_size < 2:
          return loss_dict
      with torch.no_grad():
          loss_names = []
          all_losses = []
          for k in sorted(loss_dict.keys()):
              loss_names.append(k)
              all_losses.append(loss_dict[k])
          all_losses = torch.stack(all_losses, dim=0)
          dist.reduce(all_losses, dst=0)
          if dist.get_rank() == 0:
              # only main process gets accumulated, so only divide by
              # world_size in this case
              all_losses /= world_size
          reduced_losses = {k: v for k, v in zip(loss_names, all_losses)}
      return reduced_losses
docs: Reduce the loss dictionary from all processes so that process with rank
      0 has the averaged results. Returns a dict with the same fields as
      loss_dict, after reduction.