What is torch.compile() vs TorchScript and how do you export a PyTorch model for production deployment?

PyTorch offers two main paths for production deployment beyond running the Python interpreter: TorchScript (serialises the model as a language-independent IR) and torch.compile() (JIT compiles for speed within Python). For cross-language/cross-framework deployment, ONNX export is also widely used.

import torch
import torch.nn as nn

model = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))
model.eval()   # ALWAYS call eval() before exporting

# ── 1. torch.compile() — fast Python-based inference (PyTorch 2.0+)
compiled = torch.compile(model)
with torch.no_grad():
    out = compiled(torch.randn(4, 10))

# ── 2. TorchScript trace — captures a concrete execution trace
#       Works best when control flow does NOT depend on input data
example_input = torch.randn(1, 10)
traced = torch.jit.trace(model, example_input)
torch.jit.save(traced, "model_traced.pt")

# Load and run without the original Python class
loaded_traced = torch.jit.load("model_traced.pt")
out = loaded_traced(torch.randn(4, 10))

# ── 3. TorchScript script — handles dynamic control flow
class DynamicModel(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.mean() > 0:        # data-dependent branch — trace would miss this!
            return torch.relu(x)
        return torch.tanh(x)

scripted = torch.jit.script(DynamicModel())
torch.jit.save(scripted, "model_scripted.pt")

# ── 4. ONNX export — deploy with ONNX Runtime, TensorRT, CoreML
torch.onnx.export(
    model,
    example_input,
    "model.onnx",
    input_names=["features"],
    output_names=["logits"],
    dynamic_axes={"features": {0: "batch_size"}},  # variable batch size
    opset_version=17,
)

# Inference with ONNX Runtime (no PyTorch dependency on deployment host!)
# import onnxruntime as ort
# sess = ort.InferenceSession("model.onnx")
# out  = sess.run(["logits"], {"features": x.numpy()})

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.