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

Summary:
Pull Request resolved: #1652

X-link: pytorch/pytorch#118039

X-link: pytorch/benchmark#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

torchrec/distributed/tests/test_pt2.py

@@ -127,7 +127,7 @@ def _test_kjt_input_module(
                 device = "cuda"
                 # pyre-ignore
                 aot_inductor_module = AOTIRunnerUtil.load(device, so_path)
-                aot_actual_output = aot_inductor_module(inputs)
+                aot_actual_output = aot_inductor_module(*inputs)
                 assert_close(eager_output, aot_actual_output)
 
     def test_kjt_split(self) -> None:
@@ -204,7 +204,7 @@ def kjt_to_inputs(kjt):
             )
             # pyre-ignore
             aot_inductor_module = AOTIRunnerUtil.load(device, so_path)
-            aot_inductor_module(example_inputs)
+            aot_inductor_module(*example_inputs)
 
             aot_actual_outputs = [
                 aot_inductor_module(*kjt_to_inputs(kjt)) for kjt in input_kjts[1:]