Save, Load, and Where a GPU Fits

Training takes time, and a trained model lives only in the computer's memory: when the script ends, everything it learned is gone. So you save the learned numbers to a file once, and any later script can load them and read digits right away, without training again.

Saving: model.state_dict() gives back a dictionary with one tensor per weight table and per bias, named by layer, like "0.weight". torch.save(d, "file.pt") writes it to a file. Loading: torch.load("file.pt") gives back that dictionary. It holds numbers but no layers, so you first build a new model with the same layers, then new_model.load_state_dict(d) copies each tensor into its layer. On a small model:

small = nn.Linear(3, 1)
torch.save(small.state_dict(), "small.pt")

copy = nn.Linear(3, 1)
copy.load_state_dict(torch.load("small.pt"))

Where a GPU fits: a GPU does many multiplications at once and trains big models much faster. torch.cuda.is_available() gives back True on a computer with an NVIDIA GPU, else False; Python's a if test else b picks one of two values with it. model.to(device) moves a model's weights to that device, and x.to(device) gives back a copy of the tensor x there. The model and its data must sit on the same device: that is the one change a GPU needs. This computer has no GPU, so your device will be "cpu".

The trained model and predict(model, X) (it gives back the predicted digits as a list) are already in the editor. Your code must print:

Trained model: [2, 8, 2, 6, 6, 7, 1, 9, 8, 5]
Loaded model:  [2, 8, 2, 6, 6, 7, 1, 9, 8, 5]
On cpu: [2, 8, 2, 6, 6, 7, 1, 9, 8, 5]

After your code, the check lists the four tensors in model.pt and compares the loaded model with the trained one on all 450 test images.

Your task: save the trained model's state dict to "model.pt", load it into a new model loaded and print its predictions for the first ten test images, then pick device, move loaded and the images there, and print the predictions again.

Assignment
Follow the Coding Tutorial to save the model, load it into a fresh one, and run it on the device PyTorch finds.

Hint
The new model must come from make_model() so its layers match the saved tensors. In the last print, move the images with .to(device) too, not only the model.