Revert "[aotinductor] Solves a problem where a tensor is returned more than once (#112177)"

pytorchmergebot · pytorchmergebot · commit 2bc1378d7be5 · 2023-11-06T06:20:32.000Z
This reverts commit a91baaf. Reverted #112177 on behalf of https://github.com/PaliC due to breaking internal tests (refer to internal diff) ([comment](#112177 (comment)))
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -1011,18 +1011,6 @@ def forward(self, x):
         x = torch.randn(5, device=self.device)
         self.check_model(Model(self.device), (x,))
 
-    def test_repeat_output(self):
-        class Model(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, x):
-                y = torch.sin(x)
-                return y, y
-
-        example_inputs = (torch.randn(3, 10, device=self.device),)
-        self.check_model(Model(), example_inputs)
-
 
 class AOTInductorTestABICompatibleCpu(TestCase):
     device = "cpu"
@@ -1048,8 +1036,6 @@ class AOTInductorTestABICompatibleCpu(TestCase):
         "test_freezing": TestFailure(("abi_compatible_cpu",), is_skip=True),
         "test_normal_functional": TestFailure(("abi_compatible_cpu",)),
         "test_poi_multiple_dynamic": TestFailure(("abi_compatible_cpu",)),
-        # There is a double-free issue which will be fixed in another PR
-        "test_repeat_output": TestFailure(("abi_compatible_cpu",), is_skip=True),
         "test_sdpa": TestFailure(("abi_compatible_cpu",)),
         "test_sdpa_2": TestFailure(("abi_compatible_cpu",)),
         "test_simple_dynamic": TestFailure(("abi_compatible_cpu",)),
@@ -1072,8 +1058,6 @@ class AOTInductorTestABICompatibleCuda(TestCase):
     {
         "test_dup_unbacked_sym_decl": TestFailure(("abi_compatible_cuda",)),
         "test_normal_functional": TestFailure(("abi_compatible_cuda",)),
-        # There is a double-free issue which will be fixed in another PR
-        "test_repeat_output": TestFailure(("abi_compatible_cuda",), is_skip=True),
     },
 )
 
diff --git a/torch/_export/exported_program.py b/torch/_export/exported_program.py
@@ -71,9 +71,7 @@ def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict, buf
         # Step 2: Find the all the buffers that were mutated and update them
         if node.op == "output":
             user_output_nodes = []
-            # In the case that the same node is returned multiple times,
-            # node.all_input_nodes will only iterate that node once
-            for return_node in pytree.tree_flatten(node.args)[0]:
+            for return_node in node.all_input_nodes:
                 return_node_name = return_node.name
                 # we found a param/buffer mutation
                 if return_node_name in buffers_to_mutate:
