Anybody can clarify why I keep getting cuda error when using torch.nn.parallel.DistributedDataParallel?
MWE:
if torch.cuda.device_count() > 1:
    torch.distributed.init_process_group(backend='nccl')
    model = torch.nn.parallel.DistributedDataParallel(model).to(device)

executing script:
python -m torch.distributed.launch main.py

Am I using it incorrectly or am I missing something else?
A MWE:
import os
import torch
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
torch.distributed.init_process_group(backend='nccl', rank=1, world_size=2)

Then the universe hangs completely:

image.png867×75 28.5 KB

Terminal is not released back to the user. Any ideas why?