Fix ONNX export for sequence classification head

[echarlaix]

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_indices to int32, fixing :

NOT_IMPLEMENTED : Could not find an implementation for ArgMax(13) node with name '/ArgMax'

[HuggingFaceDocBuilderDev]

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.

[echarlaix]

@ArthurZucker @Rocketknight1 let me know if that works for you

[Rocketknight1]

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.

[Rocketknight1]

(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!)

[ArthurZucker]

LGTM anyways

[Rocketknight1]

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.