How triton decide what value to return within a BaseModelOutputWithPoolingAndCrossAttentions

[sfc-gh-zhwang]

I am using triton inference server to do inference for a Bert model to do embeddings. Below is the config.pbtxt.
If I load the model locally and inference, the model would return a BaseModelOutputWithPoolingAndCrossAttentions with a last_hidden_state and pooler_output. While the triton server only returns last_hidden_state as embedding, while pooler_output is dropped. Can I know how this is decided in triton?

name: "vector_embedding"
platform: "pytorch_libtorch"
input[
    {
        name: "input_ids"
        data_type: TYPE_INT32
        dims: [1, -1]
    },
    {
        name: "attention_mask"
        data_type: TYPE_INT32
        dims: [1, -1]
    }
]
output {
    name: "embedding"
    data_type: TYPE_FP32
    dims: [1, -1, -1]
}

[dyastremsky]

Please see here: https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_configuration.md#special-conventions-for-pytorch-backend

It looks like because the naming did not follow one of the conventions, Triton assumes the input and output order matches that provided in your model. Your config only has one output, so it looks to be selecting the first output returned in your model, which I'd guess is last_hidden_state.

[zengqingfu1442]

m