ddp2.py

model = torch.nn.parallel.DistributedDataParallel(
         	model, 
			device_ids=[rank], 
			output_device=rank
        )

with model.join():
	for _ in range(5):
		for inp in inputs:
			loss = model(inp).sum()
			loss.backward()