Split trainer cli function into two functions (#1197)

talmolab · Mar 2, 2023 · 7abfac7 · 7abfac7
1 parent c636d96
commit 7abfac7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/sleap/nn/training.py b/sleap/nn/training.py
@@ -1799,8 +1799,8 @@ def visualize_example(example):
         )
 
 
-def main(args: Optional[List] = None):
-    """Create CLI for training and run."""
+def create_trainer_using_cli(args: Optional[List] = None):
+    """Create CLI for training."""
     import argparse
 
     parser = argparse.ArgumentParser()
@@ -2004,10 +2004,15 @@ def main(args: Optional[List] = None):
         test_labels=args.test_labels,
         video_search_paths=args.video_paths,
     )
-    trainer.train()
 
     return trainer
 
 
+def main(args: Optional[List] = None):
+    """Create CLI for training and run."""
+    trainer = create_trainer_using_cli(args=args)
+    trainer.train()
+
+
 if __name__ == "__main__":
     main()
diff --git a/tests/nn/test_training.py b/tests/nn/test_training.py
@@ -22,7 +22,7 @@
     TopdownConfmapsModelTrainer,
     TopDownMultiClassModelTrainer,
     Trainer,
-    main as sleap_train,
+    create_trainer_using_cli as sleap_train,
 )
 
 sleap.use_cpu_only()