Commit 61c1e160 authored by Christian Sarofeen's avatar Christian Sarofeen
Browse files

Add device sync in imagenet example.

parent a3059288
...@@ -300,6 +300,7 @@ def train(train_loader, model, criterion, optimizer, epoch): ...@@ -300,6 +300,7 @@ def train(train_loader, model, criterion, optimizer, epoch):
loss.backward() loss.backward()
optimizer.step() optimizer.step()
torch.cuda.synchronize()
# measure elapsed time # measure elapsed time
batch_time.update(time.time() - end) batch_time.update(time.time() - end)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment