Stop ct.convert from rewriting the traced model's TorchScript graph

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

@@ -437,6 +437,34 @@ def forward(self, x):
         )
         assert isinstance(model, ct.converters.mil.Program)
 
+    @staticmethod
+    def test_convert_does_not_modify_torchscript_model(tmpdir):
+        """
+        ct.convert must leave the user's TorchScript graph untouched, so the
+        traced model can still be saved and loaded afterwards (issue #2215).
+        """
+
+        class Network(torch.nn.Module):
+            def forward(self, x):
+                a, b, c = x.chunk(3)
+                return (a * b) + c
+
+        example_input = torch.rand(6, 4)
+        traced_model = torch.jit.trace(Network().eval(), example_input)
+        graph_before = str(traced_model.forward.graph)
+
+        ct.convert(
+            traced_model,
+            inputs=[ct.TensorType(name="input", shape=example_input.shape)],
+            convert_to="milinternal",
+        )
+
+        assert str(traced_model.forward.graph) == graph_before
+        path = os.path.join(tmpdir, "traced_model.pt")
+        torch.jit.save(traced_model, path)
+        loaded_model = torch.jit.load(path)
+        torch.testing.assert_close(loaded_model(example_input), traced_model(example_input))
+
     @staticmethod
     def _get_classifier_model():
         class Net(torch.nn.Module):

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

@@ -140,7 +140,9 @@ def _expand_and_optimize_ir(torchscript):
     Given a torch.jit.ScriptModule, convert it to a optimized
     torch._C.Graph and dict of model parameter's names to tensors.
     """
-    graph = torchscript.forward.graph
+    # Work on a copy: the passes below rewrite the graph in place, which would
+    # otherwise leave the user's module unable to be saved and loaded again.
+    graph = torchscript.forward.graph.copy()
 
     # From PyTorch code: Inline function and method calls.
     torch._C._jit_pass_inline(graph)