- `transformers` version: 5.3.0.dev0
- Platform: Linux-6.8.0-100-generic-x86_64-with-glibc2.39
- Python version: 3.12.12
- Huggingface_hub version: 1.5.0
- Safetensors version: 0.7.0
- Accelerate version: 1.12.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.9.1+rocm6.4 (CUDA)
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: AMD Instinct MI300X

---

from transformers import AutoModelForCausalLM
from safetensors import safe_open
from pathlib import Path
from packaging import version
import transformers

model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")

move_to_device = False  # toggle this flag.
if move_to_device:
    print("model to cuda")
    model.eval()
    model = model.to("cuda")

model.save_pretrained("opt-125m-saved")

with safe_open(Path("opt-125m-saved", "model.safetensors"), framework="pt") as f:
    checkpoint_keys = f.keys()

    for name in checkpoint_keys:
        print("checkpoint key", name)

    if move_to_device:
        if version.parse(transformers.__version__) >= version.parse("5.0"):
            # NOTE: this looks wrong / to be a regression to me.
            assert "lm_head.weight" in checkpoint_keys
        else:
            # NOTE: transformers<5 behavior looks to be the correct one to me.
            assert "lm_head.weight" not in checkpoint_keys
    else:
        assert "lm_head.weight" not in checkpoint_keys