AttributeError: 'HeadModelWithAction' object has no attribute 'save_pretrained'

[WalkerRusher]

I try to run the action-conditioned model using the following script:
accelerate launch train_gpt.py
--exp_name bair_llama_ft --output_dir log_trm --seed 0 --mixed_precision bf16
--vqgan_type ctx_vqgan
--pretrained_model_name_or_path {log directory of finetuned tokenizer}/unwrapped_model
--config_name configs/llama/config.json --load_internal_llm --action_conditioned --action_dim 4
--pretrained_transformer_path pretrained_models/ivideogpt-oxe-64-act-free/transformer
--per_device_train_batch_size 16 --gradient_accumulation_steps 1
--learning_rate 1e-4 --lr_scheduler_type cosine
--oxe_data_mixes_type bair --resolution 64 --dataloader_num_workers 16
--video_stepsize 1 --segment_length 16 --context_length 1
--use_eval_dataset --use_fvd --use_frame_metrics
--weight_decay 0.01 --llama_attn_drop 0.1 --embed_no_wd
--max_train_steps 100005

However, an error occurs:

File "/workdir/train_gpt.py", line 860, in start_train
unwrapped_model.save_pretrained(
AttributeError: 'HeadModelWithAction' object has no attribute 'save_pretrained'
raise AttributeError(f"'{type(self).name}' object has no attribute '{name}'")AttributeError
: 'HeadModelWithAction' object has no attribute 'save_pretrained'
AttributeError: 'HeadModelWithAction' object has no attribute 'save_pretrained'

[WalkerRusher]

It seems that "HeadModelWithAction" is not an instance of "Hugging Face Transformers", thus can't use the method "save_pretrained"

[Manchery]

Thanks for pointing this out! Previously, I almost saved checkpoints via Line 831 accelerator.save_state(output_dir). I will fix this issue asap.