Load the checkpoint fails

### 🐛 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_