fixes #2538

Author: tritolol

coremltools/converters/mil/frontend/torch/test/test_passes.py

@@ -18,6 +18,7 @@
     flatten_graph_input_values,
     flatten_graph_output_values,
     transform_inplace_ops,
+    remove_getattr_nodes
 )
 import coremltools as ct
 
@@ -405,3 +406,52 @@ def forward(self, x):
         y_cm = ct_model.predict({'x': x})['y']
 
         assert((y_cm == np.zeros(shape)).all())
+
+
+    @staticmethod
+    def test_remove_getattr_nodes_immediate_output():
+        graph_nodes = [
+            InternalTorchIRNode(
+                inputs=["self"],
+                attr={"name": "const_out2", "value": None},
+                outputs=["const_out2"],
+                kind="getattr",
+            ),
+            InternalTorchIRNode(
+                inputs=["self"],
+                attr={"name": "const_out1", "value": None},
+                outputs=["const_out1"],
+                kind="getattr",
+            ),
+            InternalTorchIRNode(
+                inputs=["const_out1", "const_out2"],
+                attr={"value": None},
+                outputs=["3"],
+                kind="tupleconstruct",
+            ),
+        ]
+        const2 = torch.tensor([5., 6., 7., 8.])
+        const1 = torch.tensor([1., 2., 3., 4.])
+        graph_params = {'const_out2': const2, 
+                        'const_out1': const1}
+        graph_inputs = []
+        graph_outputs = ["const_out1", "const_out2"]
+
+        graph = InternalTorchIRGraph(
+            nodes=graph_nodes,
+            params=graph_params,
+            inputs=graph_inputs,
+            outputs=graph_outputs,
+        )
+        
+        for node in graph.nodes:
+            node.parent = graph
+
+        remove_getattr_nodes(graph)
+
+        np.testing.assert_equal(graph.nodes[0].kind, "constant")
+        np.testing.assert_equal(graph.nodes[1].kind, "constant")
+        np.testing.assert_equal(graph.nodes[2].kind, "tupleconstruct")
+        np.testing.assert_allclose(graph.nodes[0].attr["value"], const2)
+        np.testing.assert_allclose(graph.nodes[1].attr["value"], const1)
+        

coremltools/converters/mil/frontend/torch/torchir_passes.py

@@ -231,10 +231,9 @@ def forward(self, x):
 
 def remove_getattr_nodes(graph: InternalTorchIRGraph) -> None:
     """
-    Remove the getattr nodes in the graph
+    Remove the getattr nodes in the graph that are not output nodes
     """
 
-    getattr_nodes = []
     new_nodes = []
 
     for node in graph.nodes:
@@ -243,16 +242,20 @@ def remove_getattr_nodes(graph: InternalTorchIRGraph) -> None:
             remove_getattr_nodes(block)
 
         if node.kind == "getattr":
-            getattr_nodes.append(node)
+            if node.name in graph.outputs:
+                # create and add new constant node
+                new_nodes.append(
+                        InternalTorchIRNode(
+                            inputs=[],
+                            outputs=node.outputs,
+                            kind="constant",
+                            name="internal_immediate_output_attr",
+                            attr={"value": node.parent.params[node.name]}
+                        )
+                    )
         else:
             new_nodes.append(node)
 
-    # check the getattr nodes not in the outputs
-    for node in getattr_nodes:
-        if node.name in graph.outputs:
-            raise RuntimeError("{} should not be in the graph outputs.".format(node.name))
-
-    # remove the getattr nodes
     graph.nodes = new_nodes