Hi @aus, that looks good to me. Just make sure that input_mask_expanded and result.hidden_state are allocated on the EXLA Backend and not the binary backend. You can print them to the terminal to confirm.
At the end, this may still be slower than the Python version for two reasons:
-
EXLA for CPU is not as fast as it should be
-
We don’t support dynamic shapes, which means you need to precompute/pad to 128. You could try passing larger inputs to both and ensure they both perform same at 128 entries or not
I believe you can call the tokenizer without a length and the size of the tensor it returns.
