3bba9a4f2dc3cecc76b105f9601a49e543f249e1,examples/data_parallel.py,,,#,29

Before Change

model = Net()
if torch.cuda.device_count() > 1:
    print("Let\"s use", torch.cuda.device_count(), "GPUs!")
    model = DataParallel(model)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

for data_list in loader:

After Change

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for data_list in loader:
    optimizer.zero_grad()
    output = model(data_list)
    y = torch.cat([data.y for data in data_list]).to(output.device)
    loss = F.nll_loss(output, y)
    loss.backward()
    optimizer.step()

Instances