Pytorch serving doesn't work for multiple inputs, multiple outputs model

Thank you for submitting a feature request. Before proceeding, please review MLflow's Issue Policy for feature requests and the MLflow Contributing Guide.

Please fill in this feature request template to ensure a timely and thorough response.

Willingness to contribute

The MLflow Community encourages new feature contributions. Would you or another member of your organization be willing to contribute an implementation of this feature (either as an MLflow Plugin or an enhancement to the MLflow code base)?

[ ] Yes. I can contribute this feature independently.
[ ] Yes. I would be willing to contribute this feature with guidance from the MLflow community.
[* ] No. I cannot contribute this feature at this time.

Proposal Summary

Currently pytorch model serving works for single input and single output model. To make it generic and flexible, it would be great if it supports multi input and multi output serving.

Motivation

What is the use case for this feature?
Why is this use case valuable to support for MLflow users in general?
Why is this use case valuable to support for your project(s) or organization?
Why is it currently difficult to achieve this use case? (please be as specific as possible about why related MLflow features and components are insufficient)

This has many use cases. In fact most of the real world models output multiple features.

What component(s), interfaces, languages, and integrations does this feature affect?

Components

[ ] area/artifacts: Artifact stores and artifact logging
[ ] area/build: Build and test infrastructure for MLflow
[ ] area/docs: MLflow documentation pages
[ ] area/examples: Example code
[ ] area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
[ *] area/models: MLmodel format, model serialization/deserialization, flavors
[ ] area/projects: MLproject format, project running backends
[ *] area/scoring: MLflow Model server, model deployment tools, Spark UDFs
[ ] area/server-infra: MLflow Tracking server backend
[ ] area/tracking: Tracking Service, tracking client APIs, autologging

Interfaces

[ ] area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
[* ] area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
[ ] area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
[ ] area/windows: Windows support

Languages

[ ] language/r: R APIs and clients
[ ] language/java: Java APIs and clients
[ ] language/new: Proposals for new client languages

Integrations

[ ] integrations/azure: Azure and Azure ML integrations
[ ] integrations/sagemaker: SageMaker integrations
[ ] integrations/databricks: Databricks integrations

Details

I am using a single input - multiple output model in this example using Pytorch lightning:

model = CustomModel(args)' 'model_dict = {'input1': np.zeros((1, 256, 256, 3))}' 'input_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 256,256,3)),])' 'output_schema = Schema([TensorSpec(np.dtype(np.float32), (-1, 2), name='out1'), TensorSpec(np.dtype(np.float32), (-1, 3), name='out2')])' 'signature = ModelSignature(inputs=input_schema, outputs=output_schema)' 'mlflow.pytorch.log_model(model,"Models", input_example=model_dict['input1'], signature=signature)

When this model is served, it throws this error:

"xxxxxxx in predict\n \"but got output of type '{}'\".format(type(preds))\nTypeError: Expected PyTorch model to output a single output tensor, but got output of type '<class 'dict'>'\n"

mlflow / mlflow