From c3f5b9aebb41988c2abb3cba443eb93015eb8b8a Mon Sep 17 00:00:00 2001
From: ntadimeti <ntadimeti@nvidia.com>
Date: Fri, 8 Jan 2021 14:00:38 -0500
Subject: [PATCH] Enable reading config files output by atacworks

---
 scripts/cmd_args.py | 6 +++---
 scripts/main.py     | 8 +++++++-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/scripts/cmd_args.py b/scripts/cmd_args.py
index aabe50ac..e25cbc9f 100644
--- a/scripts/cmd_args.py
+++ b/scripts/cmd_args.py
@@ -85,6 +85,8 @@ def add_common_options(parser):
 
     """
     # Pre-processing
+    parser.add('--config', required=False,
+               is_config_file=True, help='config file path')
     parser.add('--genome', required=True, type=str,
                help='chromosome sizes file for the genome. Sizes \
                      files for human genome 19 (hg19) and human \
@@ -142,7 +144,7 @@ def add_common_options(parser):
                help="checkpoint path to load the model from for\
                    inference or resume training")
     # dist-env args
-    parser.add('--gpu_idx', required=False, type=int,
+    parser.add('--gpu_idx', required=False, type=type_or_none_fn(int),
                help='GPU ID to use. ID can be known from nvidia-smi; \
                        preempted by --distributed which uses all \
                        available gpus ')
@@ -241,8 +243,6 @@ def add_inference_options(parser):
 
     """
     add_common_options(parser)
-    parser.add('--config', required=False,
-               is_config_file=True, help='config file path')
     parser.add('--regions', required=True, type=type_or_none_fn(str),
                help='atacworks denoising is done on whole genome by \
                      default. You can optionally specify a list of \
diff --git a/scripts/main.py b/scripts/main.py
index b4368491..a9a58dcc 100755
--- a/scripts/main.py
+++ b/scripts/main.py
@@ -584,7 +584,13 @@ def main():
     # Save config parameters
     dst_config_path = os.path.join(args.out_home,
                                    args.mode + "_config.yaml")
-    save_config(dst_config_path, args)
+    # args object is manipulated in the code, we want to save the original
+    # unchanged args object.
+    args_org = parse_args(root_dir)
+    # Delete mode from args object, as mode is registered as an argument instead of
+    # a subparser
+    del args.mode
+    save_config(dst_config_path, args_org)
 
 
 if __name__ == '__main__':