Remove reference cycle - with exceptions (#228)

* Remove reference cycle - with exceptions

* Fix for InliningInstructionTranslator

Author: anijain2305

torchdynamo/convert_frame.py

@@ -152,7 +152,6 @@ def transform(instructions, code_options):
             )
             tracer.run()
             output = tracer.output
-            output.cleanup()
             assert output.output_instructions
             instructions[:] = output.output_instructions
             code_options.update(output.code_options)

torchdynamo/output_graph.py

@@ -104,6 +104,10 @@ def restore_graphstate(self, state):
         # FX deepcopy doesn't work for a partially created graph, so just remove new nodes
         for node in reversed(list(self.graph.nodes)):
             if node not in graph_nodes:
+                # Erasing node alone does not remove the meta information
+                # So, remove the help tensor explicitly
+                if "example_value" in node.meta:
+                    del node.meta["example_value"]
                 self.graph.erase_node(node)
 
     def count_calls(self):
@@ -374,3 +378,7 @@ def cleanup(self):
         # Cleanup graphargs
         for graph_arg in self.graphargs:
             graph_arg.erase()
+
+        for node in self.graph.nodes:
+            if "example_value" in node.meta:
+                del node.meta["example_value"]

torchdynamo/symbolic_convert.py

@@ -314,6 +314,14 @@ def run(self):
                 f"{self.lineno} {typestr(e)}\n"
             )
             raise
+        finally:
+            # Cleanup the outputGraph to delete the held tensors. We perform the
+            # cleanup only for InstructionTranslator and not
+            # InliningInstructionTranslator. The InliningInstructionTranslator
+            # mutates the output object and is restored to original state if
+            # there was an exception.
+            if isinstance(self, InstructionTranslator):
+                self.output.cleanup()
 
     def push(self, val):
         assert val is None or isinstance(