In our case, we need to return different graphs to run according to different application scenarios. In order to avoid the dynamo cache returning the same graph every time, we choose to let the custom compiler return an empty graph and save the differently optimized graphs, and then call the specific graphs outside compiler like this:
# 1. In compiler:
def __call__(original_gm: torch.fx.GraphModule, example_inputs):
...
compiler.fwd_gm = original_gm
return get_none_module(original_gm, self.variable_pool, self.alias_analysis)
# 2. run real graph module
output = compiler.run_gm(compiler.fwd_gm)
However, it just return some tuple-type values rather than the expected dict-type values. Do you know how to make the return value as expected the same as running compiled model directly?