Skip to content

Instantly share code, notes, and snippets.

@AFAgarap
Last active March 6, 2024 03:07
Show Gist options
  • Save AFAgarap/3129def5ac7f8471f5429768bfa01ae0 to your computer and use it in GitHub Desktop.
Save AFAgarap/3129def5ac7f8471f5429768bfa01ae0 to your computer and use it in GitHub Desktop.
Loading MNIST dataset and creating a torch.utils.data.DataLoader object for it.
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
train_dataset = torchvision.datasets.MNIST(
root="~/torch_datasets", train=True, transform=transform, download=True
)
test_dataset = torchvision.datasets.MNIST(
root="~/torch_datasets", train=False, transform=transform, download=True
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True
)
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=32, shuffle=False, num_workers=4
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment