Fix ONNX export for sequence classification head #36332
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@ArthurZucker @Rocketknight1 let me know if that works for you |
Rocketknight1
left a comment
Rocketknight1
left a comment
There was a problem hiding this comment.
Ah, I wrote this code, but I didn't realize this would cause an issue for ONNX! The change is fine, since the output should be identical regardless.
We could also consider updating the line
pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]to
pooled_logits = logits[torch.arange(batch_size, device=logits.device, dtype=torch.int32), last_non_pad_token]I don't think it's strictly necessary, but if we don't do this then the indexing will use an int64 tensor for one dimension and an int32 tensor for another dimension. That works fine on Torch right now, but it seems like the kind of thing that might break in other frameworks or when the model is exported.
|
(If you don't want to do that, though, we can just merge this PR as-is. The change I suggested is just future-proofing rather than strictly necessary!) |
|
Yeah, let's just merge this as-is! If the int64/int32 inputs thing becomes a problem when exporting later we can fix it then. |
* set dtype to int32 * fix style
4 participants
Since #35911, some models with a sequence classification head will be exported to ONNX with an ArgMax operator in int64, making the resulting model incompatible with ORT. Here we set
token_indicesto int32, fixing :