Skip to content

Instantly share code, notes, and snippets.

@tkshnkmr

tkshnkmr/check_dataloader.py

Created September 28, 2019 14:22
Show Gist options
  • Save tkshnkmr/bac5c28246e051b3166ea1c69e798db8 to your computer and use it in GitHub Desktop.
Save tkshnkmr/bac5c28246e051b3166ea1c69e798db8 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
check_dataloader.py
# 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, targets in data_loader:
imgs = list(img.to(device) for img in imgs)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
print(targets)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment