Update train_lora.py

zejunwang1 · Jul 19, 2023 · d70a9ac · d70a9ac
1 parent 93febd1
commit d70a9ac
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/train_lora.py b/train_lora.py
@@ -105,6 +105,10 @@ def train():
     # cast all non INT8 parameters to fp32
     model = prepare_model_for_kbit_training(model, 
         use_gradient_checkpointing=training_args.gradient_checkpointing)
+
+    if training_args.gradient_checkpointing:
+        model.enable_input_require_grads()
+        model.gradient_checkpointing_enable()
 
     # Get our peft model and print the number of trainable parameters
     checkpoint_dir = training_args.resume_from_checkpoint