def _get_test_model_2():
        """
        Base model have the same weights, while the weights in submodule are different.
        """

        class SubModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear1 = torch.nn.Linear(28 * 28 * 8, 5, bias=False)

            def forward(self, x):
                return self.linear1(x)


        class TestModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.conv1 = torch.nn.Conv2d(1, 8, 5, padding="same", bias=False)
                self.bn1 = torch.nn.BatchNorm2d(8)
                self.linear1 = None

            def forward(self, x):
                x = self.conv1(x)
                x = self.bn1(x)
                x = self.linear1(torch.flatten(x))
                return x


        example_input = torch.rand(1, 1, 28, 28)
        model = TestModel().eval()

        submodule_1 = SubModel().eval()
        model.linear1 = submodule_1
        trace_1 = torch.jit.trace(model, example_input)

        submodule_2 = SubModel().eval()
        model.linear1 = submodule_2
        trace_2 = torch.jit.trace(model, example_input)

        return trace_1, trace_2

        Base model have the same weights, while the weights in submodule are different.
        
coremltools/test/ml_program/test_utils.py