diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py
index 00fb7a5..8f8acfe 100644
--- a/scripts/train_merge_code.py
+++ b/scripts/train_merge_code.py
@@ -51,7 +51,7 @@ ap.add_argument("--deepk", type=int, default=0,
 ap.add_argument("--bptt", type=int, default=0,
                 help="truncated BPTT: grads only through last N iterations")
 ap.add_argument("--lr", type=float, default=1e-3)
-ARGS = ap.parse_args()
+ARGS, _ = ap.parse_known_args()
 SEED = ARGS.seed
 LR = ARGS.lr
 SUFFIX = ((f"_s{SEED}" if SEED else "")