1 files changed, 5 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index 6c57f4b..7f5fb49 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -12,10 +12,12 @@ import torch.utils.checkpoint
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
-from slugify import slugify
 from timm.models import create_model
 import transformers
 
+import numpy as np
+from slugify import slugify
+
 from util.files import load_config, load_embeddings_from_dir
 from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, add_placeholder_tokens, get_models
@@ -852,6 +854,7 @@ def main():
     )
 
     data_generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    data_npgenerator = np.random.default_rng(args.seed)
 
     def run(i: int, placeholder_tokens: list[str], initializer_tokens: list[str], num_vectors: Union[int, list[int]], data_template: str):
         placeholder_token_ids, initializer_token_ids = add_placeholder_tokens(
@@ -894,6 +897,7 @@ def main():
             filter=partial(keyword_filter, filter_tokens, args.collection, args.exclude_collections),
             dtype=weight_dtype,
             generator=data_generator,
+            npgenerator=data_npgenerator,
         )
         datamodule.setup()