Description
🐛 Describe the bug
Here's a simple repro.
Legacy onnx.export works.
Both dynamo_export and export.export + dynamo_export fail:
import torch
model = torch.nn.GLU().eval().cuda()
sample_input = (torch.randn(8, 4).cuda(),)
# This works fine
torch.onnx.export(model, sample_input, "glu_legacy.onnx")
model = torch.export.export(model, sample_input)
# This line has to be uncommented as well for dynamo_export to succeed
# Otherwise, the following error is produced by dynamo_export:
# torch.onnx._internal.diagnostics.infra.context.RuntimeErrorWithDiagnostic: Unsupported FX nodes: {'call_function': ['aten\
.glu.default']}.
# model = model.run_decompositions()
options = torch.onnx.ExportOptions(dynamic_shapes=True)
out = torch.onnx.dynamo_export(model,
*sample_input,
export_options=options
)
out.save("glu.onnx")
### Versions
PyTorch nightly 05/09
cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang