def ddp_setup(rank: int, world_size: int):
    """
    Set environment variables which are used for initializing distributed
    process group for :py:class:`DistributedDataParallel`.

    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    """
    _os.environ["MASTER_ADDR"] = "localhost"
    _os.environ["MASTER_PORT"] = "12355"
    _os.environ["WORLD_SIZE"] = f"{world_size}"
    _os.environ["RANK"] = f"{rank}"
    _os.environ["LOCAL_RANK"] = f"{rank}"
    _torch.cuda.set_device(f"cuda:{rank}")
    _dist.init_process_group("nccl", rank=rank, world_size=world_size)

    Set environment variables which are used for initializing distributed
    process group for :py:class:`DistributedDataParallel`.

    Args:
        rank: Unique identifier of each process
        world_size: Total number of processes
    
coremltools/optimize/torch/_utils/dist_utils.py