add log dir

calico · May 20, 2024 · 2bf5178 · 2bf5178
1 parent e26ece9
commit 2bf5178
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/baskerville/scripts/hound_train.py b/src/baskerville/scripts/hound_train.py
@@ -56,6 +56,12 @@ def main():
         default="train_out",
         help="Output directory [Default: %(default)s]",
     )
+    parser.add_argument(
+        "-l",
+        "--log_dir",
+        default=None,
+        help="Tensorboard log directory [Default: %(default)s]",
+    )
     parser.add_argument(
         "--restore",
         default=None,
@@ -150,7 +156,7 @@ def main():
 
         # initialize trainer
         seqnn_trainer = trainer.Trainer(
-            params_train, train_data, eval_data, args.out_dir
+            params_train, train_data, eval_data, args.out_dir, args.log_dir
         )
 
         # compile model
@@ -182,6 +188,7 @@ def main():
                 train_data,
                 eval_data,
                 args.out_dir,
+                args.log_dir,
                 strategy,
                 params_train["num_gpu"],
                 args.keras_fit,