Set TOKENIZERS_PARALLELISM to false in train.py (#120)

Avoids the warning about "The current process just got forked. Disabling parallelism to avoid deadlocks"
unitaryai · Jan 3, 2025 · c77352b · c77352b
1 parent 50d16db
commit c77352b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -1,9 +1,12 @@
+# flake8: noqa: E402
 import argparse
 import json
 import os
 
-import pytorch_lightning as pl
+# Disable tokenizer parallelism to avoid deadlocks due to forking
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+import pytorch_lightning as pl
 import src.data_loaders as module_data
 import torch
 from pytorch_lightning.callbacks import ModelCheckpoint