Skip to content

Instantly share code, notes, and snippets.

@tkshnkmr
Last active September 28, 2019 14:24
Show Gist options
  • Save tkshnkmr/81a761d9e7c3a55d029d00b125efeb23 to your computer and use it in GitHub Desktop.
Save tkshnkmr/81a761d9e7c3a55d029d00b125efeb23 to your computer and use it in GitHub Desktop.
# select device (whether GPU or CPU)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# DataLoader is iterable over Dataset
for imgs, annotations in data_loader:
imgs = list(img.to(device) for img in imgs)
annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
print(annotations)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment