tqdm_tutorial_mnist_load.py

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.optim import Adam
from torch import nn

transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
dataset = datasets.MNIST('./mnist_data', train=True, download=True,
                         transform=transform)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Perceptron().to(device)

optimizer = Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
batch_size = 64

train_loader = DataLoader(dataset,
                          batch_size=batch_size,
                          num_workers=1,
                          pin_memory=True,
                          shuffle=True,
                          drop_last=True)