From f66a2032bfedba9c6d949ed0cec946b295566b03 Mon Sep 17 00:00:00 2001 From: Om Singhal Date: Sun, 13 Sep 2026 15:52:39 -0400 Subject: [PATCH] Stop ct.convert from rewriting the traced model's TorchScript graph --- .../torch/test/test_torch_conversion_api.py | 28 +++++++++++++++++++ .../mil/frontend/torch/torchscript_utils.py | 4 ++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py b/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py index 07760adc9..829290334 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py +++ b/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): diff --git a/coremltools/converters/mil/frontend/torch/torchscript_utils.py b/coremltools/converters/mil/frontend/torch/torchscript_utils.py index 4712501a5..a9b5ccddf 100644 --- a/coremltools/converters/mil/frontend/torch/torchscript_utils.py +++ b/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)