Exceptions in Torch Export

torch.compile
1

Exception are currently not supported in torch.export.export. See the example below:

import os
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        raise RuntimeError("Exception")

with torch.no_grad():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Model().to(device=device)
    example_inputs=(torch.randn(8, 10, device=device),)
    torch.export.export(model, example_inputs)

which results in the following (abbreviated) error:

Unsupported: call_function BuiltinVariable(RuntimeError) [ConstantVariable(str)] {}

from user code:
   File "/tmp/main.py", line 9, in forward
    raise RuntimeError("Exception")

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

However, exception are an idiomatic way to report errors and bail-out early. For example, I count 278 uses of exceptions in diffusers models folder.

Is there a way to handle exceptions in exported code or what would be an idiomatic way to bail out early in programs? I can think of using torch.cond and separate functions in two paths and return error codes if needed, but that doesn’t seem to be very ergonomic.