Skip to content

Instantly share code, notes, and snippets.

@Siyeong-Lee
Created May 13, 2019 07:41
Show Gist options
  • Select an option

  • Save Siyeong-Lee/bf228ddc8c595b3089e823538404ff2b to your computer and use it in GitHub Desktop.

Select an option

Save Siyeong-Lee/bf228ddc8c595b3089e823538404ff2b to your computer and use it in GitHub Desktop.
import torch
import torch.nn as nn
from parallel import DataParallelModel, DataParallelCriterion
model = BERT(args)
model = DataParallelModel(model)
model.cuda()
criterion = nn.NLLLoss()
criterion = DataParallelCriterion(criterion)
...
for i, (inputs, labels) in enumerate(trainloader):
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment