how to understand the local_rank and if..else in this code?

Viewed 76

I do not know the two ways of setting the device and what the local rank refers to. Can anybody explain this code to me?

if args.local_rank == -1:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
else:
    torch.distributed.init_process_group(backend='nccl')
    torch.cuda.set_device(args.local_rank)
    device = torch.device('cuda', args.local_rank)
0 Answers
Related