Using oneDNN with 🤗 models

Hi.

I have been trying to understand the performance benefits of OneDNN for CPU inference (reference).

To this end, I am trying benchmark a pre-trained model with and without OneDNN optimizations. Inference with the traced model seems to be breaking with:

RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.FloatTensor instead (while checking arguments for embedding)

Any suggestions to resolve it?

Here’s my Colab Notebook: Google Colab.