PyTorch TorchScript

TorchScript converts a PyTorch model from Python into an intermediate representation that runs without a Python interpreter. This lets you deploy models in production environments where Python is unavailable — C++ applications, mobile devices, embedded systems — while keeping full PyTorch compatibility during development.

Why TorchScript Exists

Python PyTorch model:
  ✓ Flexible, easy to debug, great for research
  ✗ Requires Python runtime to execute
  ✗ Python's Global Interpreter Lock limits multi-threading
  ✗ Cannot run on mobile apps or C++ servers

TorchScript model:
  ✓ Runs without Python
  ✓ Portable to C++, mobile, edge devices
  ✓ Can be optimized and serialized to a single file
  ✓ Thread-safe for concurrent inference

Two Ways to Create TorchScript

Method 1: torch.jit.script
  → Analyzes your Python code and converts it
  → Works with control flow (if/else, loops)
  → Requires type annotations in some cases

Method 2: torch.jit.trace
  → Runs the model with example input, records operations
  → Simpler but misses data-dependent control flow
  → Best for straightforward feed-forward models

Method 1: torch.jit.script

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 16)
        self.fc2 = nn.Linear(16, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleNet()

# Convert to TorchScript via scripting
scripted_model = torch.jit.script(model)

x   = torch.rand(8, 4)
out = scripted_model(x)
print(out.shape)   # torch.Size([8, 2])

# Save to disk
scripted_model.save('simple_net_scripted.pt')

# Load it back (no class definition needed)
loaded = torch.jit.load('simple_net_scripted.pt')
print(loaded(x).shape)   # torch.Size([8, 2])

Method 2: torch.jit.trace

model = SimpleNet()

# Provide an example input — tracing records the operations
example_input  = torch.rand(1, 4)
traced_model   = torch.jit.trace(model, example_input)

print(traced_model(torch.rand(8, 4)).shape)   # torch.Size([8, 2])

# Save
traced_model.save('simple_net_traced.pt')

Script vs Trace — When Each Breaks

Model with data-dependent control flow:

  def forward(self, x):
      if x.sum() > 0:     ← depends on actual data value
          return self.fc1(x)
      else:
          return self.fc2(x)

  → torch.jit.trace records only ONE branch (the one taken by example input)
    Any input that goes down the other branch gives WRONG results silently.
  → torch.jit.script handles both branches correctly.

Rule:
  Use script for models with if/else or loops that depend on input values.
  Use trace for simple sequential models with no conditional logic.

Inspecting TorchScript Code

scripted = torch.jit.script(SimpleNet())
print(scripted.code)
# Shows the TorchScript IR (intermediate representation):
# def forward(self, x: Tensor) -> Tensor:
#   x = torch.relu(self.fc1.forward(x))
#   return self.fc2.forward(x)

TorchScript with Control Flow

class BranchNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc_pos = nn.Linear(4, 2)
        self.fc_neg = nn.Linear(4, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.mean() > 0:
            return self.fc_pos(x)
        else:
            return self.fc_neg(x)

scripted = torch.jit.script(BranchNet())   # handles both branches

x_pos = torch.rand(3, 4)     # mean > 0 → uses fc_pos
x_neg = -torch.rand(3, 4)    # mean < 0 → uses fc_neg

print(scripted(x_pos).shape)
print(scripted(x_neg).shape)

Running TorchScript Models from C++

After saving a scripted model, load it from a C++ application using LibTorch. This is how PyTorch models are deployed in high-performance servers, games, and robotics.

C++ code (conceptual):
  #include <torch/script.h>

  torch::jit::script::Module model;
  model = torch::jit::load("simple_net_scripted.pt");
  model.eval();

  auto input  = torch::rand({8, 4});
  auto output = model.forward({input}).toTensor();

Optimizing for Inference

scripted = torch.jit.script(model)

# Optimize for mobile or CPU inference
optimized = torch.jit.optimize_for_inference(scripted)

x   = torch.rand(1, 4)
out = optimized(x)
print(out)

TorchScript on Mobile

TorchScript is the bridge to PyTorch Mobile (iOS and Android). After scripting and optimizing, serialize the model and bundle it with your app. The model runs fully on-device — no server call required.

from torch.utils.mobile_optimizer import optimize_for_mobile

scripted    = torch.jit.script(model)
mobile_opt  = optimize_for_mobile(scripted)
mobile_opt._save_for_lite_interpreter('model_mobile.ptl')

Summary

TorchScript converts PyTorch models to a Python-independent format that runs anywhere — C++ servers, mobile apps, and edge hardware. Use torch.jit.script for models with conditional logic or loops; use torch.jit.trace for simple sequential models. Save with .save() and load with torch.jit.load() — no class definition needed at load time. Inspect the generated code with scripted.code. Use optimize_for_mobile before deploying to iOS or Android. TorchScript is the recommended path for all PyTorch production deployments that need to operate outside a Python environment.