chore: remove mp from training too

AnswerDotAI · Jan 16, 2024 · e24fe44 · e24fe44
1 parent 066218a
commit e24fe44
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/ragatouille/models/colbert.py b/ragatouille/models/colbert.py
@@ -336,6 +336,8 @@ def _batch_search(self, query: list[str], k: int):
     def train(self, data_dir, training_config: ColBERTConfig):
         training_config = ColBERTConfig.from_existing(self.config, training_config)
         training_config.nway = 2
+        if training_config.nranks < 2:
+            training_config.avoid_fork_if_possible = True
         with Run().context(self.run_config):
             trainer = Trainer(
                 triples=str(data_dir / "triples.train.colbert.jsonl"),