I see. Another question is regarding this piece of code:
    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    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(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

The model has been designated to a GPU and also wrapped by DDP. But when we feed in data as in this line
outputs = ddp_model(torch.randn(20, 10))

Shouldn’t we use torch.randn(20, 10).to(rank) instead?