Load the checkpoint fails

[AfraAmini]

🐛 Describe the bug

I am training a model using PPO with 8 unfrozen layers. I have the checkpoint saved under hf_model directory. When I attempt to load the checkpoint I get the error that layers are not initialized from the checkpoint and there are extra keys (base_model. and frozen_head.) that were not matched with any keys in the GPT model. Here is how I attempted to load the model:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("ckpts/checkpoint_6400/hf_model")

Which trlX version are you using?

No response

Additional system and package information

No response