ValueError: The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models.

[WinstonLiyt]

Here is my code:

model = AutoModelForCausalLMWithValueHead.from_pretrained(
        config.model_name,
        load_in_4bit=True,
        device_map="auto",
        peft_config=lora_config,
        quantization_config=BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
    )

Terminal output (Error): File "predict_module/tuning_lm_with_rl.py", line 154, in tuning_lm_with_rl model = AutoModelForCausalLMWithValueHead.from_pretrained( File "trl/models/modeling_base.py", line 308, in from_pretrained model.post_init(state_dict=state_dict) File "trl/models/modeling_value_head.py", line 238, in post_init raise ValueError( ValueError: The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models.

Can anyone help me? TKS!!