Update

1 files changed, 5 insertions, 1 deletions

diff --git a/train_lora.py b/train_lora.py
index 91bda5c..d5dde02 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -13,9 +13,11 @@ from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import LoggerType, set_seed
 from peft import LoraConfig, LoraModel
-from slugify import slugify
 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
@@ -873,6 +875,7 @@ def main():
     )
 
     data_generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    data_npgenerator = np.random.default_rng(args.seed)
 
     create_datamodule = partial(
         VlpnDataModule,
@@ -893,6 +896,7 @@ def main():
         valid_set_pad=args.valid_set_pad,
         dtype=weight_dtype,
         generator=data_generator,
+        npgenerator=data_npgenerator,
     )
 
     create_lr_scheduler = partial(