def _average_pool_func(model_output, attention_mask):
        """A pooling function that averages the hidden states of the model."""
        last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
A pooling function that averages the hidden states of the model.
src/marqo/inference/native_inference/embedding_models/hugging_face_model.py