Fix AIMv2 checkpoint/config compatibility with custom HF code

#7
by tewth - opened

This PR makes the uploaded checkpoint consistent with the custom modeling_aimv2.py implementation shipped in the repository.

  • remap tensor names in model.safetensors to match AIMv2Model.state_dict()
  • update config.json so architectures and model_type align with the shipped config/model classes

Without these changes, loading via AutoModel.from_pretrained(..., trust_remote_code=True) initializes most model weights from scratch instead of loading the published checkpoint.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment