Medusa convert cpt doesn't work with safetensors version

[nivibilla]

System Info

N/A

Who can help?

No response

Information

• [x] The official example scripts
• [ ] My own modified scripts

Tasks

• [x] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• [ ] My own task or dataset (give details below)

Reproduction

try converting a medusa checkpoint using text-generation-inference/Mistral-7B-Instruct-v0.2-medusa

Expected behavior

Should load from safe tensors

actual behavior

script expects a "medusa_lm_head.pt" in folder

additional notes

n/a

[rakib-hasan]

Thanks for reporting this. @nivibilla You are correct. The script needs to be updated to be able to load safe tensors as well.

[Tushar-ml]

Opened a PR #1535 for this

[Tushar-ml]

@ncomly-nvidia kindly review this PR

[Tushar-ml]

This has been merged, we can close this issue