fix: batch size

MinishLab · Feb 3, 2025 · e1d53ac · e1d53ac
1 parent 261a9b4
commit e1d53ac
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py
@@ -129,7 +129,9 @@ def fit(
         )
 
         if batch_size is None:
-            batch_size = max(min(32, len(train_texts) // 10), 512)
+            # Set to a multiple of 32
+            base_number = int(min(max(1, (len(train_texts) / 30) // 32), 16))
+            batch_size = int(base_number * 32)
             logger.info("Batch size automatically set to %d.", batch_size)
 
         logger.info("Preparing train dataset.")