1 files changed, 3 insertions, 1 deletions

diff --git a/train_lora.py b/train_lora.py
index 4bbc64e..0d8ee23 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -865,6 +865,8 @@ def main():
         max_grad_norm=args.max_grad_norm,
     )
 
+    data_generator = torch.Generator(device="cpu").manual_seed(args.seed)
+
     create_datamodule = partial(
         VlpnDataModule,
         data_file=args.train_data_file,
@@ -882,8 +884,8 @@ def main():
         valid_set_size=args.valid_set_size,
         train_set_pad=args.train_set_pad,
         valid_set_pad=args.valid_set_pad,
-        seed=args.seed,
         dtype=weight_dtype,
+        generator=data_generator,
     )
 
     create_lr_scheduler = partial(