Save differentiable model head on CPU

[tomaarsen]

Supersedes #327

Hello!

Pull Request overiew

• Save differentiable model head on CPU
• Move differentiable model heads to the right device after loading

Details

See #327 for more information:

If we train a differentiable head on a CUDA machine, the head checkpoint cannot be loaded via Joblib on CPU machine (Serialization error).
To address this issue, I modified the code to store the head checkpoint with CPU mapping.
When we load the head checkpoint, we map it to the target device.

Thanks @karter-liner for providing the PR. I've simply applied roughly their changes on top of the upcoming v1.0.0-pre branch.

• Tom Aarsen

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.