What I understood is that you are trying to switch off gradient checkpointing, which is fine if you have enough GPU memory.

TO achieve this, before you define the training configs of the Trainer class, use the prepare_model_for_kbit_training() method.

The way you do this is below:

model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)