Fix a bug in the torch._export.aot_load API (#2128)

Summary:
X-link: pytorch/torchrec#1652

X-link: pytorch/pytorch#118039

Pull Request resolved: #2128

tree_flatten_spec should use args instead of *args

clone of pytorch/pytorch#117948 but with some fbcode specific changes

Reviewed By: angelayi

Differential Revision: D52982401

fbshipit-source-id: 503326c8b4d7316315153a4ab3873f6c8bcac9fb

userbenchmark/dynamo/dynamobench/common.py

@@ -1149,7 +1149,7 @@ def export_aot_inductor(model, example_inputs, device):
 
     def opt_aot_inductor(_, example_inputs, collect_outputs=False):
         example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
-        return optimized(example_args, example_kwargs)
+        return optimized(*example_args, **example_kwargs)
 
     return opt_aot_inductor