From 33a6b2cb0596e58014ad829fbb78164505493304 Mon Sep 17 00:00:00 2001
From: Han Zhu <zhuhan568@outlook.com>
Date: Tue, 21 May 2024 23:38:37 -0400
Subject: [PATCH] Support batch size adjustment

---
 main.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index b68e944..3933bcc 100644
--- a/main.py
+++ b/main.py
@@ -14,6 +14,7 @@
     parser.add_argument("-F", nargs="+", type=str, help="List of input mrc files", required=True)
     parser.add_argument("-G", type=str, help="GPU ID to use for prediction", required=False, default="")
     parser.add_argument("-J", type=str, help="Job name / output folder name", required=True)
+    parser.add_argument("-B", type=int, help="Batch size to use", required=False, default=4)
 
     args = parser.parse_args()
 
@@ -56,7 +57,7 @@
                 f"-F={mrc_file}",
                 "--contour=0",
                 f"--gpu={args.G}",
-                f"--batch_size=4",
+                f"--batch_size={args.B}",
                 f"--prediction_only",
                 f"--resolution=8.0",
                 f"--output={curr_out_dir}",