I am a bit confused about averaging gradients in distributed data-parallel. It seems there are two examples from the PyTorch documentation that are different. In one example, you create the model and just pass it to the GPU available then create a separate function to average gradients.
def average_gradients(model):
    """ Gradient averaging. """
    size = float(dist.get_world_size())
    for param in model.parameters():
        dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM)
        param.grad.data /= size

and its executed as follows.
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        average_gradients(model)
        optimizer.step()


The other approach I have seen doesn’t create a separate function and just calls DPP.
    model = ToyModel().cuda(device_ids[0])
    ddp_model = DDP(model, device_ids)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_ids[0])
    loss_fn(outputs, labels).backward()
    optimizer.step()

I would like to know whats the difference between the two approaches and which one should one use for distributed training in a HPC cluster. I specifically want to use two nodes, each with 4 GPUs.