1 files changed, 5 insertions, 3 deletions

diff --git a/data/textual_inversion/csv.py b/data/textual_inversion/csv.py
index e082511..3ac57df 100644
--- a/data/textual_inversion/csv.py
+++ b/data/textual_inversion/csv.py
@@ -19,7 +19,8 @@ class CSVDataModule(pl.LightningDataModule):
                  interpolation="bicubic",
                  placeholder_token="*",
                  center_crop=False,
-                 valid_set_size=None):
+                 valid_set_size=None,
+                 generator=None):
         super().__init__()
 
         self.data_file = Path(data_file)
@@ -35,6 +36,7 @@ class CSVDataModule(pl.LightningDataModule):
         self.center_crop = center_crop
         self.interpolation = interpolation
         self.valid_set_size = valid_set_size
+        self.generator = generator
 
         self.batch_size = batch_size
 
@@ -48,10 +50,10 @@ class CSVDataModule(pl.LightningDataModule):
     def setup(self, stage=None):
         valid_set_size = int(len(self.data_full) * 0.2)
         if self.valid_set_size:
-            valid_set_size = math.min(valid_set_size, self.valid_set_size)
+            valid_set_size = min(valid_set_size, self.valid_set_size)
         train_set_size = len(self.data_full) - valid_set_size
 
-        self.data_train, self.data_val = random_split(self.data_full, [train_set_size, valid_set_size])
+        self.data_train, self.data_val = random_split(self.data_full, [train_set_size, valid_set_size], self.generator)
 
         train_dataset = CSVDataset(self.data_train, self.tokenizer, size=self.size, repeats=self.repeats, interpolation=self.interpolation,
                                    placeholder_token=self.placeholder_token, center_crop=self.center_crop)