examples/demo_torch_multi_gpu.py/setup
def setup(current_gpu_index, num_gpu):
    # Device setup
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "56492"
    device = torch.device("cuda:{}".format(current_gpu_index))
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        world_size=num_gpu,
        rank=current_gpu_index,
    )
    torch.cuda.set_device(device)