fix OOM problem

RayJue · Apr 8, 2023 · 384968b · 384968b
1 parent 3d6f0e9
commit 384968b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/finetune_chatglm.py b/finetune_chatglm.py
@@ -40,6 +40,8 @@ def main():
 
     # Training
     if training_args.do_train:
+        model.gradient_checkpointing_enable()
+        model.enable_input_require_grads()
         trainer.train()
         model.save_pretrained(training_args.output_dir)