Set TOKENIZERS_PARALLELISM to false in train.py

Avoids the warning about "The current process just got forked. Disabling parallelism to avoid deadlocks"
unitaryai · Jan 1, 2025 · ea654a4 · ea654a4
1 parent 6a9b738
commit ea654a4
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