feat: Updated sweepai/config/server.py

sweepai · Mar 15, 2024 · f622f2c · f622f2c
1 parent bbd4c04
commit f622f2c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/sweepai/config/server.py b/sweepai/config/server.py
@@ -30,9 +30,11 @@
     "SENTENCE_TRANSFORMERS_MODEL",
     "sentence-transformers/all-MiniLM-L6-v2",  # "all-mpnet-base-v2"
 )
-BATCH_SIZE = int(
+from sweepai.utils.multiprocessing_utils import safe_multiprocessing_setup
+
+BATCH_SIZE = safe_multiprocessing_setup(int(
     os.environ.get("BATCH_SIZE", 256)
-)  # Tune this to 32 for sentence-transformers/all-MiniLM-L6-v2 on CPU
+))  # Tune this to 32 for sentence-transformers/all-MiniLM-L6-v2 on CPU
 
 TEST_BOT_NAME = "sweep-nightly[bot]"
 ENV = os.environ.get("ENV", "dev")