Changes

The code is very similar to a single GPU implementation, except for the ''.to('cuda:x')'' calls, where ''cuda:0'' and ''cuda:1'' are each their own GPU.

model = ToyModel()

loss_fn = nn.MSELoss()

loss_fn(outputs, labels).backward()

optimizer.step()

The backward() and torch.optim will automatically take care of gradients as if the model is on one GPU. You only need to make sure that the labels are on the same device as the outputs when calling the loss function.