Load the checkpoint fails #600
Description
🐛 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