def test_input_output_description(self):
        """
        When using save_multifunction to produce a model, we should respect
        the original model description in the original model.
        """

        def check_i_o_spec(model_path: str) -> None:
            spec = load_spec(model_path)
            model_desc = spec.description

            assert len(model_desc.functions) == 2

            for idx in [0, 1]:
                assert (
                    model_desc.functions[idx].input[0].type.imageType.colorSpace
                    == proto.FeatureTypes_pb2.ImageFeatureType.BGR
                )
                assert (
                    model_desc.functions[idx].output[0].type.imageType.colorSpace
                    == proto.FeatureTypes_pb2.ImageFeatureType.RGB
                )

        # source model with i/o with ImageType
        class Model(torch.nn.Module):
            def forward(self, x):
                return x + 5.0

        example_input = torch.randint(0, 100, (1, 3, 10, 20), dtype=torch.float32)
        model = torch.jit.trace(Model().eval(), example_input)
        mlmodel = ct.convert(
            model,
            inputs=[ct.ImageType(shape=(1, 3, 10, 20), color_layout=ct.colorlayout.BGR)],
            outputs=[ct.ImageType(color_layout=ct.colorlayout.RGB)],
            minimum_deployment_target=ct.target.iOS16,
        )
        package_path = tempfile.mkdtemp(suffix=".mlpackage")
        mlmodel.save(package_path)

        # multifunction model should have the same i/o description
        model_path = self._multifunction_model_from_single_function(package_path)
        check_i_o_spec(model_path)

        # construct another multifunction model with an existing multifunction model,
        # the i/o description should still be the same
        model_path_2 = self._multifunction_model_from_multifunction_model(model_path)
        check_i_o_spec(model_path_2)

        # cleanup
        shutil.rmtree(package_path)
        shutil.rmtree(model_path)
        shutil.rmtree(model_path_2)

        When using save_multifunction to produce a model, we should respect
        the original model description in the original model.
        
coremltools/test/ml_program/test_utils.py