dynamo_capture_subgraphs
is a bit odd function. You need to do something like tvm.IRModule({"main": mod["subgraph_0"]})
.
Note that PT 2.0 import story in Relax is currently not good primarily because PT 2.0 compile situation is not mature at the moment, especially wrt whole-graph export that many of TVM users care about but unfortunately not the first priority for Meta.
For one data point, in web-stable-diffusion the authors had to manually reimplement SD UNet just so that it can be traced by FX. The authors are the same people who developed the FX frontend in Relax.