ddp.py

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")
model = MyModel().cuda()
model = DDP(model, device_ids=[local_rank])

for batch in dataloader:
    optimizer.zero_grad()
    outputs = model(batch)
    loss = criterion(outputs, targets)
    loss.backward()  # gradients synchronized across processes automatically
    optimizer.step()