ONNX Export Pitfalls: 7 PyTorch → Production Gotchas

The Export Worked. The Inference Failed.

Your PyTorch model exports to ONNX without errors. The file loads in ONNX Runtime. Then inference produces NaN outputs, silently wrong predictions, or crashes with cryptic shape errors.

I’ve debugged this cycle enough times to spot the patterns. ONNX export doesn’t fail loudly — it fails in production, after you’ve already committed to the deployment strategy. Most tutorials stop at torch.onnx.export() succeeding, but that’s where the real problems start.

Here are seven gotchas I’ve hit moving PyTorch models to ONNX, complete with the specific errors, fixes, and when each one bites you.

1. Dynamic Axes Break When You Don’t Expect Them

You specify dynamic_axes={'input': {0: 'batch'}} during export. The model runs fine with batch size 1, 8, 32 during validation. Then you deploy with batch size 7 and get a shape mismatch error buried six ops deep in the graph.

import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
import numpy as np

class ConvClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.pool(x)
        # Bug here: reshape assumes batch size is known at export time
        x = x.view(-1, 16)  # This bakes in assumptions
        return self.fc(x)

model = ConvClassifier()
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "classifier.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}}
)

# Validation with same batch size: works
sess = ort.InferenceSession("classifier.onnx")
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = sess.run(None, {'input': test_input})
print(f"Batch 1 output shape: {output[0].shape}")  # (1, 10) ✓

# Different batch size: fails or gives wrong results
test_input_batch8 = np.random.randn(8, 3, 224, 224).astype(np.float32)
try:
    output = sess.run(None, {'input': test_input_batch8})
    print(f"Batch 8 output shape: {output[0].shape}")
except Exception as e:
    print(f"Error: {e}")

The problem is view(-1, 16). ONNX tries to infer shapes at export time, and operations like view or reshape that depend on computed dimensions can hardcode batch size.

Fix: use flatten(start_dim=1) or explicitly handle the batch dimension:

class ConvClassifierFixed(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.pool(x)
        x = torch.flatten(x, start_dim=1)  # Explicitly preserves batch dim
        return self.fc(x)

Test this with multiple batch sizes BEFORE you deploy. Export with batch=1, then validate with 1, 4, 7, 16, 32. If any fail, your dynamic axes aren’t actually dynamic.

2. In-Place Operations Silently Corrupt Gradients (Even in Inference)

You’re doing inference only, so gradients don’t matter, right? Wrong. ONNX export traces the computational graph, and in-place ops like x += residual or x.relu_() can create aliasing that breaks graph optimization passes.

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        residual = x
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x += residual  # In-place add — ONNX export warning
        x = self.relu(x)
        return x

model = ResidualBlock(64)
model.eval()
dummy = torch.randn(1, 64, 32, 32)

torch.onnx.export(model, dummy, "residual_inplace.onnx", opset_version=17)
# UserWarning: ONNX export mode is set to TrainingMode.EVAL, but operator 'aten::add_' is set to TrainingMode.TRAINING.

You’ll get a warning, but the export completes. The exported model might work, or it might produce subtly wrong outputs after ONNX Runtime applies graph optimizations. I’ve seen cases where the residual connection got optimized away entirely because the optimizer thought the in-place add was safe to remove.

Fix: never use in-place operations in models you plan to export.

class ResidualBlockFixed(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
        self.relu = nn.ReLU()

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.relu(out)
        out = self.conv2(out)
        out = out + residual  # Out-of-place add
        out = self.relu(out)
        return out

This adds one tensor allocation per forward pass in PyTorch, but it’s negligible compared to the conv ops. And your ONNX export won’t silently break.

3. Opset Version Compatibility: The Export Succeeds, Runtime Fails

You export with opset_version=17 (PyTorch 2.1 default). Your production server runs ONNX Runtime 1.12, which supports up to opset 16. The model loads without error, then crashes on the first op that uses opset 17 features.

# Export with opset 17
torch.onnx.export(
    model,
    dummy_input,
    "model_opset17.onnx",
    opset_version=17
)

# Load in ONNX Runtime 1.12 (supports opset ≤ 16)
sess = ort.InferenceSession("model_opset17.onnx")
# onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : 
# Node (Reshape_23) Op (Reshape) [ShapeInferenceError] Opset 17 not supported

The error message doesn’t always say “opset version mismatch” — it often manifests as “unsupported operator” or “shape inference error” on a specific node.

Check your runtime’s supported opset before export:

print(f"ONNX Runtime version: {ort.__version__}")
print(f"Supported opsets: {ort.get_available_providers()}")
# For opset compatibility, check onnxruntime docs or:
# https://github.com/microsoft/onnxruntime/blob/main/docs/Versioning.md

I default to opset 13 for maximum compatibility (supported since ONNX Runtime 1.8, released mid-2021). Only use newer opsets if you specifically need features like SequenceMap or OptionalHasElement.

torch.onnx.export(
    model,
    dummy_input,
    "model_opset13.onnx",
    opset_version=13  # Safe for most production environments
)

4. Control Flow Gets Unrolled (Loops and Conditionals)

You have a model with dynamic control flow — a loop that runs for n iterations where n depends on input, or a conditional branch based on tensor values. PyTorch handles this fine. ONNX export tries to trace it and either fails or unrolls the loop for the specific input size you used during export.

class DynamicLoopModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x, num_iterations):
        # num_iterations is a Python int, not a tensor
        for _ in range(num_iterations):
            x = self.fc(x)
            x = torch.relu(x)
        return x

model = DynamicLoopModel()
model.eval()
dummy_input = torch.randn(1, 10)

# Export with num_iterations=3
torch.onnx.export(
    model,
    (dummy_input, 3),
    "dynamic_loop.onnx",
    input_names=['input', 'num_iterations'],
    dynamic_axes={'input': {0: 'batch'}}
)
# The exported model hardcodes 3 iterations. Changing num_iterations at runtime does nothing.

ONNX does support control flow via Loop and If ops (opset ≥ 13), but PyTorch’s tracing-based export doesn’t emit them by default. You need to use torch.jit.script instead of torch.jit.trace, then export the scripted model:

class DynamicLoopModelScripted(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x, num_iterations: int):
        for _ in range(num_iterations):
            x = self.fc(x)
            x = torch.relu(x)
        return x

model_scripted = torch.jit.script(DynamicLoopModelScripted())
torch.onnx.export(
    model_scripted,
    (dummy_input, torch.tensor(3)),
    "dynamic_loop_scripted.onnx",
    opset_version=13
)

This works, but ONNX Runtime’s support for Loop ops varies by execution provider. CPU execution provider handles it fine. TensorRT and CoreML might not. If your deployment target doesn’t support control flow ops, you’ll need to refactor your model to avoid dynamic loops.

My rule: if the loop count is fixed (e.g., always 3 transformer layers), unrolling is fine. If it genuinely needs to be dynamic, reconsider whether ONNX is the right export format. TorchScript might be a better fit.

5. Custom Ops: The Black Box Problem

You implemented a custom CUDA kernel or a PyTorch autograd function. Export fails with “RuntimeError: ONNX export failed: Couldn’t export operator my_custom_op”.

class CustomOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        # Custom logic: element-wise square + 1
        return x * x + 1

    @staticmethod
    def backward(ctx, grad_output):
        # Gradients (not used in inference, but required for autograd)
        return grad_output * 2 * ctx.saved_tensors[0]

class ModelWithCustomOp(nn.Module):
    def forward(self, x):
        return CustomOp.apply(x)

model = ModelWithCustomOp()
dummy = torch.randn(1, 10)

torch.onnx.export(model, dummy, "custom_op.onnx")
# RuntimeError: ONNX export failed: Couldn't export operator aten::custom_op

Two options:

1. Register a symbolic function that maps your custom op to ONNX ops:

from torch.onnx import register_custom_op_symbolic

def custom_op_symbolic(g, x):
    # Map to ONNX ops: x^2 + 1
    x_squared = g.op("Mul", x, x)
    one = g.op("Constant", value_t=torch.tensor([1.0]))
    return g.op("Add", x_squared, one)

register_custom_op_symbolic('::custom_op', custom_op_symbolic, 13)

This works if your custom op can be expressed as a composition of standard ONNX ops. If it’s a genuinely novel operation (like a custom attention mechanism with fused CUDA kernels), you’re out of luck.

2. Refactor to avoid the custom op. Replace it with equivalent PyTorch primitives:

class ModelWithoutCustomOp(nn.Module):
    def forward(self, x):
        return x * x + 1  # Same logic, no custom autograd function

Most “custom ops” I’ve seen in production are actually just compositions of standard ops wrapped for convenience. If you can unwrap them, do it before export.

6. Hardcoded Constants vs. Initializers

You have a normalization layer that subtracts the mean $\mu = [0.485, 0.456, 0.406]$ and divides by std $\sigma = [0.229, 0.224, 0.225]$. You define these as Python lists or torch.tensor() inside forward(). ONNX export bakes them into the graph as Constant nodes, but some runtimes don’t optimize constant folding well, leading to redundant compute.

class NormalizeBad(nn.Module):
    def forward(self, x):
        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
        return (x - mean) / std

model = NormalizeBad()
dummy = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy, "normalize_bad.onnx")

# Check exported graph
import onnx
onnx_model = onnx.load("normalize_bad.onnx")
print(f"Num constant nodes: {sum(1 for n in onnx_model.graph.node if n.op_type == 'Constant')}")
# 6 constant nodes (mean values, std values, shape constants for reshape)

Better: register them as buffers so they become initializers in the ONNX graph:

class NormalizeGood(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward(self, x):
        return (x - self.mean) / self.std

model = NormalizeGood()
torch.onnx.export(model, dummy, "normalize_good.onnx")

onnx_model = onnx.load("normalize_good.onnx")
print(f"Num initializers: {len(onnx_model.graph.initializer)}")
# 2 initializers (mean, std) — cleaner graph, better optimized

This matters more than you’d think. I’ve debugged a model where constant folding wasn’t happening, and the runtime was recalculating the same normalization constants for every batch. Moving them to initializers cut inference time by 8%.

7. NaN Outputs: Mismatched Default Values

Your PyTorch model uses F.interpolate(..., mode='bilinear'). The default align_corners is None in PyTorch (interpreted as False for some modes). ONNX’s Resize op has different defaults. The exported model produces slightly different outputs, and if you’re chaining ops that are sensitive to tiny differences (like softmax with very small logits), you get NaN gradients or extreme values.

import torch.nn.functional as F

class UpsampleModel(nn.Module):
    def forward(self, x):
        # align_corners not specified — PyTorch default varies by mode
        return F.interpolate(x, scale_factor=2, mode='bilinear')

model = UpsampleModel()
dummy = torch.randn(1, 3, 32, 32)

torch.onnx.export(model, dummy, "upsample.onnx")

# Compare PyTorch vs ONNX Runtime outputs
with torch.no_grad():
    pytorch_out = model(dummy).numpy()

sess = ort.InferenceSession("upsample.onnx")
onnx_out = sess.run(None, {'input': dummy.numpy()})[0]

max_diff = np.abs(pytorch_out - onnx_out).max()
print(f"Max difference: {max_diff:.6f}")
# Max difference: 0.000347 — small, but compounds through deeper models

Fix: always specify ambiguous arguments explicitly.

class UpsampleModelFixed(nn.Module):
    def forward(self, x):
        return F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)

This also applies to torch.nn.functional.grid_sample (padding_mode, align_corners), torch.topk (sorted), and torch.nn.functional.pad (mode, value). If an argument has a default, set it explicitly before export.

Validation Pipeline: Catch These Before Deployment

Exporting to ONNX isn’t one command. It’s a pipeline:

1. Export with verbose logging: torch.onnx.export(..., verbose=True) shows warnings you’d otherwise miss.
2. Validate the graph: onnx.checker.check_model("model.onnx") catches malformed graphs.
3. Numerical parity check: run the same inputs through PyTorch and ONNX Runtime, assert outputs match within tolerance:

def validate_onnx(pytorch_model, onnx_path, dummy_input, rtol=1e-3, atol=1e-5):
    pytorch_model.eval()
    with torch.no_grad():
        pytorch_out = pytorch_model(dummy_input).numpy()

    sess = ort.InferenceSession(onnx_path)
    onnx_out = sess.run(None, {'input': dummy_input.numpy()})[0]

    np.testing.assert_allclose(pytorch_out, onnx_out, rtol=rtol, atol=atol)
    print(f"✓ Validation passed (max diff: {np.abs(pytorch_out - onnx_out).max():.2e})")

validate_onnx(model, "model.onnx", torch.randn(1, 3, 224, 224))

4. Test multiple batch sizes if using dynamic axes.
5. Benchmark on target hardware. ONNX Runtime CPU vs GPU vs TensorRT can have different bugs.

I run this validation in CI before merging any model architecture changes. It’s caught regressions where a “harmless” refactor broke ONNX export.

When ONNX Isn’t Worth It

Not every model should export to ONNX. If you’re hitting multiple issues from this list, consider alternatives:

• TorchScript: supports dynamic control flow, custom ops, easier debugging. Downside: locked into PyTorch runtime.
• TensorRT directly: if you’re targeting NVIDIA GPUs, TensorRT’s PyTorch parser can be more forgiving than ONNX export.
• TorchServe or Triton with PyTorch backend: skip export entirely, serve the .pt checkpoint. You lose some optimization opportunities but gain compatibility.

I’ve worked on projects where we spent two weeks debugging ONNX export only to realize TorchScript deployment was simpler and faster. ONNX shines when you need true framework interop (PyTorch training → C++/mobile/web inference) or when the target runtime (ONNX Runtime, CoreML, TFLite via ONNX) has hardware-specific optimizations you can’t replicate elsewhere.

For a typical FastAPI deployment on AWS EC2, I’d stick with FastAPI Model Serving: 5 Steps to 50ms Inference and skip ONNX unless latency profiling proves the export is worth the debugging time.

FAQ

Q: Can I export a model that uses torch.compile()?

Yes, but you need to export the original uncompiled model. torch.compile() applies graph optimizations that are PyTorch-runtime-specific and won’t translate to ONNX. Call torch.onnx.export() on the base nn.Module before compiling, or keep a reference to the uncompiled version.

Q: Why does my ONNX model run slower than PyTorch?

Usually because ONNX Runtime isn’t using the optimal execution provider for your hardware. Check available providers with ort.get_available_providers(). If you have a GPU but only see CPUExecutionProvider, reinstall onnxruntime-gpu. Also verify you’re not accidentally running in debug mode — ONNX Runtime’s default logging is verbose and slow.

Q: How do I debug which node in the ONNX graph is producing NaN?

Use Netron (https://netron.app) to visualize the graph and identify node names. Then run inference with intermediate outputs enabled:

sess = ort.InferenceSession("model.onnx")
output_names = [n.name for n in sess.get_outputs()]
intermediate_names = [n.name for n in sess.get_graph().node][:10]  # First 10 nodes
outputs = sess.run(output_names + intermediate_names, {'input': test_input})
for name, out in zip(output_names + intermediate_names, outputs):
    print(f"{name}: NaN={np.isnan(out).any()}, range=[{out.min():.2e}, {out.max():.2e}]")

This shows you exactly where the NaN first appears. Common culprits: division by zero in normalization layers, log of negative values, softmax overflow.

Pick Your Pain Points

Use ONNX when you need cross-platform deployment or hardware-specific runtimes that beat PyTorch’s performance. Accept that you’ll spend time validating numerical parity and debugging shape mismatches.

Skip it if your deployment is Python-only, your model has heavy control flow, or you’re prototyping and iteration speed matters more than 10ms latency gains.

I lean toward ONNX for production CV models (ResNet-style architectures, object detection) where the graph is mostly static and TensorRT gives real speedups. For LLMs or RL policies with dynamic compute graphs, I stick with TorchScript or native PyTorch serving.

The biggest lesson: validate on real production data and batch sizes before committing. The errors that matter don’t show up in unit tests with torch.randn(1, 3, 224, 224). They show up at 3am when a user uploads a 4000×3000 image and your dynamic axes weren’t actually dynamic. When debugging that at midnight, Dark Chocolate Espresso Beans are the only dependency I trust unconditionally.