Loading a Saved PyTorch Model

After a model has been fine-tuned, the values it has learned exist only while the program is running. When the program ends, those values are gone unless they were written to a file. Saving is therefore the final step of the fine-tuning workflow: the fine-tuned model is stored on disk so it can be reused in future projects, and later loaded back into a program.

This article follows the saving and loading path one line at a time. As you read, pay attention to which value each line produces before the next line uses it. Most save and load problems come from a value produced on one line that does not match what a later line expects.

The saving work happens in one call: torch.save(model.state_dict(), file_name). The first argument is model.state_dict(), which gives the model's state dictionary. The second argument is file_name, a string that names the file on disk. The model is saved as a .pth file.

The file name is computed on the line before the save, so the name is fixed before torch.save runs. The save call does not decide where the output goes on its own; the file name alone determines the destination. This is why a wrong file name produces a file with the wrong name rather than a missing save.

The saved object is the state dictionary, not the whole program. Calling model.state_dict() gives a dictionary-style collection of the model's stored values. Each entry in this collection is associated with a name, and that name points to one stored value. Loading later depends on those names, which is why the saved structure must be recovered in the same form.

The file name is not typed by hand. It is generated from the chosen model name. The code takes CHOOSE_MODEL, removes every space, opening parenthesis, and closing parenthesis using re.sub with the pattern [ ()], and then appends the suffix -sft.pth. The result is a name such as gpt2-medium355M-sft.pth.

Loading runs in the opposite direction to saving. The file is read by torch.load, the resulting state dictionary is passed to load_state_dict, and the model receives the saved values. The order matters: load_state_dict is called on a model object, and the argument it receives comes from torch.load.

Read the load line from the inside out. torch.load(...) is evaluated first and produces a state dictionary. That dictionary is then passed as the argument to model.load_state_dict, which applies it to the model. The model object is the one that changes; the file itself is only read.

When a save or load fails, trace the file name first, since both steps must use the same generated string. Then walk the sequence in order: the name is built, the state dictionary is passed to torch.save, the file is read by torch.load, and the result is applied by load_state_dict. The first step where the values or names stop matching is the one to inspect.