import os
import torch
import torch.distributed as dist
import datetime

def main():
    dist.init_process_group("nccl",timeout=datetime.timedelta(seconds=10))
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    print(f"Rank {dist.get_rank()} initialized on GPU {local_rank}")

    x = torch.ones(10).cuda()
    dist.all_reduce(x)
    print(f"Rank {dist.get_rank()} result: {x[0].item()}")

if __name__ == "__main__":
    main()