Read Handwritten Digits

The file digits.pt holds 1,797 handwritten digits, each an image of 8x8 = 64 pixels. The starting code reads them with torch.load, divides the pixels by 16 so they go from 0 to 1, keeps the right digit of each image in y, and hands them out in batches of 64. Your job is the network that reads them.

nn.Sequential(...) takes layers and gives back one model that runs them in order: each layer's output is the next layer's input. nn.ReLU() is the activation between the layers, where your from-scratch network had sigmoid: it gives back its input with every negative number turned into 0. A small example with other sizes:

tiny = nn.Sequential(nn.Linear(3, 4), nn.ReLU(), nn.Linear(4, 2))
print(tiny(torch.rand(5, 3)).shape)                    # torch.Size([5, 2]): 5 rows, 2 numbers each
print(nn.ReLU()(torch.tensor([-2.0, 0.5])).tolist())   # [0.0, 0.5]

To pick one of 10 digits, the last layer gives 10 scores, one per digit; the highest score is the model's answer. nn.CrossEntropyLoss() is the loss for choosing between classes: it turns the scores into probabilities and is small when the right digit gets a high one. It takes the right digit itself, like 5, as the answer. torch.optim.Adam(model.parameters(), lr=0.01) is an optimizer like SGD that adjusts the step size of every weight by itself, so it usually learns faster; the training loop uses it exactly like SGD, and PyTorch still works out every slope for you.

The loop prints the loss on all the images every 5 epochs. Your program should print:

epoch 0 loss: 0.708
epoch 5 loss: 0.114
epoch 10 loss: 0.069
epoch 15 loss: 0.06
epoch 20 loss: 0.029

The check then reads the first ten images (9 right: it reads the 5 as a 9), finds 0.996 of the training images right, and counts 2,410 numbers the model learns.

Your task: below the loader, build model with nn.Sequential: a Linear layer from 64 to 32, a ReLU, and a Linear layer from 32 to 10. Then store nn.CrossEntropyLoss() in loss_fn and an Adam optimizer with lr=0.01 in optimizer.

Assignment
Follow the Coding Tutorial to build the digits network, its loss and its optimizer.

Hint
Each layer inside nn.Sequential( ... ) sits on its own line, four spaces in, with a comma after it. The sizes must chain: the first layer's 32 outputs are the second layer's 32 inputs.