ValueError: EsmForTokenClassification does not support gradient checkpointing

[suresh-pokharel]

Hi Amelie, I am looking to finetune ESM model on my small downstream task. I am following your article https://huggingface.co/blog/AmelieSchreiber/esm2-ptm.

I am getting the following error on gradient_checkpointing_enable() function that you might already have a solution/suggestions.

    model.gradient_checkpointing_enable() # SP commented
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sureshp/anaconda3/envs/qlora/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1631, in gradient_checkpointing_enable
    raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
ValueError: EsmForTokenClassification does not support gradient checkpointing.

Thanks.