diff --git a/seq2seq/run_seq2seq.py b/seq2seq/run_seq2seq.py
index 9749f4a0..4ec3ea9f 100644
--- a/seq2seq/run_seq2seq.py
+++ b/seq2seq/run_seq2seq.py
@@ -122,8 +122,8 @@ def main() -> None:
         num_beams=data_training_args.num_beams,
         num_beam_groups=data_training_args.num_beam_groups,
         diversity_penalty=data_training_args.diversity_penalty,
-        gradient_checkpointing=model_args.gradient_checkpointing,
-        use_cache=not model_args.gradient_checkpointing,
+        gradient_checkpointing=training_args.gradient_checkpointing,
+        use_cache=not training_args.gradient_checkpointing,
     )
 
     # Initialize tokenizer
diff --git a/seq2seq/utils/args.py b/seq2seq/utils/args.py
index 91100214..25673faf 100644
--- a/seq2seq/utils/args.py
+++ b/seq2seq/utils/args.py
@@ -37,7 +37,3 @@ class ModelArguments:
             "with private models)."
         },
     )
-    gradient_checkpointing: bool = field(
-        default=False,
-        metadata={"help": "Will enable gradient / activation checkpointing"},
-    )