[Feature]: load_weights function in JambaForSequenceClassification

[afrenkai]

🚀 The feature, motivation and pitch

Working on my first contribution, and wanted to start with something light. I was reading #10860 which had a TODO for explicitly casting weights to float 32, but at the moment it just casts self.scores to float, and does not affect weights. I would like to submit a little 2 liner to have weights cast explicitly to float32 and score to float32. If I am misunderstanding something, please let me know. As far as I am aware, there are no open PRs or issues addressing this small change.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.