Here is a toy model that contains two linear layers. Each linear layer is designed to run a separate GPU.

import torch

x = self.relu(self.net1(x.to('cuda:0')))

return self.net2(x.to('cuda:1'))

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(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.