# Inference on models with custom head

**URL:** <https://discuss.huggingface.co/t/inference-on-models-with-custom-head/137621>\
**Category:** 🤗Optimum\
**Created:** [January 27, 2025, 9:57pm UTC](https://discuss.huggingface.co/t/inference-on-models-with-custom-head/137621 "2025-01-27T21:57:50Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![kaushu42](https://avatars.discourse-cdn.com/v4/letter/k/a8b319/32.png) [@kaushu42](https://discuss.huggingface.co/u/kaushu42)\
**Post date:** [January 27, 2025, 9:57pm UTC](https://discuss.huggingface.co/t/inference-on-models-with-custom-head/137621/1 "2025-01-27T21:57:50Z")

</div>

Hi,  
I trained a BERT model with a custom classifier head. I tried to load is using the `ORTModelForSequenceClassification`, but it did not load the proper classifier heads. It just loaded the default heads.

After several hours of digging through the code for optimum, I was able to get my code to work using this:

```auto
from optimum.onnxruntime import ORTModelForSequenceClassification
from optimum.exporters.tasks import TasksManager
TasksManager.infer_library_from_model = lambda *args, **kwargs: "transformers"
TasksManager.get_model_class_for_task = lambda *args, **kwargs: BertWithCustomHead

onnx_model = ORTModelForSequenceClassification.from_pretrained("./my-classifier", export=True, task="text-classification")

```

It gets the job done, but I feel it’s kind of hacky. Does anyone know if there is there a better way to do this?

---

<div class="post-metadata">

**Author:** ![echarlaix](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/echarlaix/32/21754_2.png) [@echarlaix](https://discuss.huggingface.co/u/echarlaix)\
**Post date:** [January 28, 2025, 11:18am UTC](https://discuss.huggingface.co/t/inference-on-models-with-custom-head/137621/2 "2025-01-28T11:18:08Z")

</div>

Hi @kaushu42, when loading your model with `ORTModelForSequenceClassification`, the original model will be loaded before conversion with `AutoModelForSequenceClassification`. In your case, you should first load your model and then use [`onnx_export_from_model`](https://huggingface.co/docs/optimum/main/en/exporters/onnx/package_reference/export#optimum.exporters.onnx.onnx_export_from_model) to export it to ONNX