🎓 A1Learner
🔍
0 of 71 lessons completed

PyTorch Saving Scripted Models

A scripted model saves to one file with torch.jit.save. The file loads with torch.jit.load, and no class code is needed at load time.

Save and load in one program

The loaded copy answers on its own after the round trip. The input 3 comes back as 7.

import torch

class Doubler(torch.nn.Module):
    def forward(self, x):
        return x * 2 + 1

scripted = torch.jit.script(Doubler())
torch.jit.save(scripted, "snippet_model.pt")
loaded = torch.jit.load("snippet_model.pt")
print(loaded(torch.tensor([3.0])).item())

The saved file is real

The save writes a real file with real size on disk. The check reads the file size after saving.

import os
import torch

class Doubler(torch.nn.Module):
    def forward(self, x):
        return x * 2 + 1

scripted = torch.jit.script(Doubler())
torch.jit.save(scripted, "snippet_model.pt")
print(os.path.getsize("snippet_model.pt") > 1000)

Loading after the class is gone

The class name is deleted before the load here. The loaded model still runs, because the file carries its own code.

import torch

class Doubler(torch.nn.Module):
    def forward(self, x):
        return x * 2 + 1

scripted = torch.jit.script(Doubler())
torch.jit.save(scripted, "snippet_model.pt")
del Doubler
loaded = torch.jit.load("snippet_model.pt")
print(loaded(torch.tensor([3.0])).item())

Example

import subprocess
import sys
import torch

class Doubler(torch.nn.Module):
    def forward(self, x):
        return x * 2 + 1

torch.manual_seed(0)
model = Doubler()
scripted = torch.jit.script(model)
torch.jit.save(scripted, "saved_model.pt")

loader = (
    "import torch\n"
    "loaded = torch.jit.load('saved_model.pt')\n"
    "out = loaded(torch.tensor([3.0]))\n"
    "print(f'Loaded prediction: {out.item():.1f}')\n"
)
with open("loader.py", "w") as f:
    f.write(loader)

result = subprocess.run(
    [sys.executable, "-P", "loader.py"],
    capture_output=True, text=True)
line = result.stdout.strip()
print(line)
own = scripted(torch.tensor([3.0])).item()
print(f"Original prediction: {own:.1f}")
print(f"Match: {'Loaded prediction: ' + format(own, '.1f') == line}")

Key points

  • Use torch.jit.save(scripted, path) to write a scripted model to one file
  • Use torch.jit.load(path) to read the file back as a runnable model
  • The file holds the compiled code and the weights together
  • A fresh process loads the file without importing the model class

In this example, a small scripted model is saved to saved_model.pt. The program then writes a tiny loader program with no class definition in it. A subprocess runs that loader with a fresh Python interpreter. The loader prints its own prediction of 7.0 for the input 3. The original model prints the same 7.0, and the match line reads True.

Note

Torch 2.14 marks jit.save as deprecated, and its warning points to torch.export

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does a program need to load a saved scripted model?