56
edits
Changes
→Parallelization Methods
Here is a toy model that contains two linear layers. Each linear layer is designed to run a separate GPU.
import torch import torch.nn as nn import torch.optim as optim
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.